Fix Megatron installer runtime integration

Author: Kovbo

src/art/megatron/setup.sh

@@ -9,4 +9,4 @@ apt-get install -y libcudnn9-headers-cuda-12 ninja-build
 
 # Python dependencies are installed through art_megatron_install so
 # downstream repos can reuse the same source of truth for versions and VCS pins.
-uv run python -m art_megatron_install
+uv run --no-sync python -m art_megatron_install

src/art_megatron_install.py

@@ -47,6 +47,11 @@
     "transformer-engine-torch",
     "megatron-core",
 )
+_BRIDGE_OVERRIDE_LINES = (
+    "megatron-core==0.16.1\n",
+    "transformer-engine==2.11.0\n",
+    "quack-kernels==0.2.5\n",
+)
 
 
 @dataclass(frozen=True)
@@ -201,16 +206,23 @@ def install_megatron_dependencies(*, dry_run: bool = False) -> None:
 
     bridge_requirement = requirements.pop("megatron-bridge", None)
     if bridge_requirement is not None:
+        override_path = "/tmp/art-megatron-bridge-override.txt"
+        if not dry_run:
+            with open(override_path, "w", encoding="utf-8") as override_file:
+                override_file.writelines(_BRIDGE_OVERRIDE_LINES)
         _run(
             [
                 _uv_executable(),
                 "pip",
                 "install",
-                "--no-deps",
+                "--override",
+                override_path,
                 _format_requirement(bridge_requirement),
             ],
             dry_run=dry_run,
         )
+        if not dry_run and os.path.exists(override_path):
+            os.remove(override_path)
 
     if dry_run or _package_installed("pynvml"):
         _run([_uv_executable(), "pip", "uninstall", "pynvml"], dry_run=dry_run)