Megatron LoRA correctness: align distributed semantics with Megatron and validate TP/EP/ETP/DP against an oracle (#619)

* megatron: integrate lora grad sync with finalize_model_grads

* megatron: harden sharded lora merge validation

* tests: add megatron lora oracle correctness harness

* Minor typing changes

* megatron: extend LoRA grad-sync semantics across tp/expert-tp

* megatron: add MoE routing replay core and unit tests

* megatron runtime/service: wire routing replay into training jobs

* oracle worker/trace: capture forward traces and emit replay bundles

* oracle harness/tests: refactor suite and add oracle-replay parity flow

* typing: clear blocking ty errors in oracle replay and LoRA paths

* megatron: reduce oracle variance with sequence grad accumulation

Use per-step micro-accumulation over multiple packed sequences so updates are less sensitive to sparse expert token assignment. Also make backend progress accounting accumulation-aware.

* megatron lora: fix TP/EP export participation rules

Correct LoRA shard export behavior so non-zero TP ranks in EP/ETP topologies contribute when required, while still filtering replicated-only entries.

* oracle trace: canonicalize MoE outputs across arbitrary topologies

Move normalization logic into ForwardTraceCapture so saved traces are canonicalized toward world-size-1 semantics (expert row identity/order and ETP fc1 layout).

* oracle harness: stabilize scoring and expand sensitivity mutations

Rework oracle pass/fail evaluation with per-phase functions, layer-averaged metrics, deterministic init, expanded sensitivity mutations, and smaller Adam epsilon for tiny-gradient regimes.

* oracle tests: write suite output tables to log files

Redirect suite stdout/stderr into local correctness/sensitivity logs and make skip/report messaging point to those artifacts instead of terminal output.

* Add correct data parallelism.

* Fix per-token DP normalization in Megatron training

* Expand the oracle harness for DP correctness checks

* Clean up type errors in Megatron correctness changes

* Testing harness was working, but real training surfaced a few errors, mostly fixed.

* Cut over Megatron LoRA to QuACK

* Del held packed tensors so dir can be removed.

Plus small typing changes.

* Fuse LoRA scale into QuACK grouped GEMM

* Avoid grad_out copy in QuACK LoRA backward

* Fuse MoE FC1 gate and up LoRA paths

* Tune QuACK low-rank tiles and rank contract

* Inline FC1 QuACK dual call

* Revert unnecessary python 3.12 requirement.

* Create lora without instantiating full model by using meta device.

- Fix routing replay for torch 2.10.0.

* Update Megatron dependencies for transformers v5 change.

* Update megatron tests for new lora kernel and avg grads across experts for stability.

* Limit max build jobs when building the uv cache.

* Fix CI uv cache build robustness

* Tune CI uv cache build concurrency

* Fix CI Apex cache contract

Author: FurtherAI

.github/workflows/prek.yml

@@ -13,6 +13,9 @@ env:
   CI_PYTHON_MM: "3.11"
   CI_UV_CACHE_RELEASE_TAG: "prek-uv-cache"
   CI_UV_CACHE_ASSET_PREFIX: "prek-uv-cache"
+  CI_APEX_PARALLEL_BUILD: "8"
+  CI_APEX_NVCC_THREADS: "1"
+  CI_UV_BUILD_SLOTS: "2"
   UV_CACHE_DIR: "/root/.cache/uv"
   UV_LINK_MODE: "copy"
   TORCH_CUDA_ARCH_LIST: "8.0"
@@ -34,7 +37,9 @@ jobs:
             --pyproject pyproject.toml \
             --uv-lock uv.lock \
             --base-image "${CI_BASE_IMAGE}" \
-            --python-mm "${CI_PYTHON_MM}")"
+            --python-mm "${CI_PYTHON_MM}" \
+            --ci-apex-parallel-build "${CI_APEX_PARALLEL_BUILD}" \
+            --ci-apex-nvcc-threads "${CI_APEX_NVCC_THREADS}")"
           echo "fingerprint=${fp}" >> "${GITHUB_OUTPUT}"
           echo "Expected uv cache fingerprint: ${fp}"
 
@@ -198,6 +203,13 @@ jobs:
 
       - name: Install dependencies (with all optional extras for complete type checking)
         run: |
+          original_pyproject="$(mktemp)"
+          cp pyproject.toml "${original_pyproject}"
+          cleanup() {
+            mv "${original_pyproject}" pyproject.toml
+          }
+          trap cleanup EXIT
+
           py_mm="$(python -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')"
           cudnn_path="${GITHUB_WORKSPACE}/.venv/lib/python${py_mm}/site-packages/nvidia/cudnn"
           export CUDNN_PATH="${cudnn_path}"
@@ -207,13 +219,22 @@ jobs:
           export CPLUS_INCLUDE_PATH="${CUDNN_INCLUDE_PATH}${CPLUS_INCLUDE_PATH:+:${CPLUS_INCLUDE_PATH}}"
           export LIBRARY_PATH="${CUDNN_LIBRARY_PATH}${LIBRARY_PATH:+:${LIBRARY_PATH}}"
           export LD_LIBRARY_PATH="${CUDNN_LIBRARY_PATH}${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"
+          export UV_CONCURRENT_BUILDS="${CI_UV_BUILD_SLOTS}"
+          export CMAKE_BUILD_PARALLEL_LEVEL="${CI_APEX_PARALLEL_BUILD}"
+          export MAX_JOBS="${CI_APEX_PARALLEL_BUILD}"
+          export NINJAFLAGS="-j${CI_APEX_PARALLEL_BUILD}"
+          python3 scripts/ci/apply_ci_uv_build_overrides.py \
+            --pyproject pyproject.toml \
+            --apex-parallel-build "${CI_APEX_PARALLEL_BUILD}" \
+            --apex-nvcc-threads "${CI_APEX_NVCC_THREADS}"
+          echo "CI uv build overrides: APEX_PARALLEL_BUILD=${CI_APEX_PARALLEL_BUILD}, NVCC_APPEND_FLAGS=--threads ${CI_APEX_NVCC_THREADS}, UV_CONCURRENT_BUILDS=${CI_UV_BUILD_SLOTS}"
           uv --version
           uv sync --all-extras --group dev --frozen
 
       - name: Run prek hooks (lint, format, typecheck, uv.lock, tests)
         run: |
-          uv run prek run --all-files
+          uv run --no-sync prek run --all-files
 
       - name: Run unit tests (via prek)
         run: |
-          uv run prek run pytest
+          uv run --no-sync prek run pytest