Merge upstream main

Signed-off-by: Philip Petrakian <ppetrakian@nvidia.com>

Author: Phlip79

.github/workflows/cicd-main.yml

@@ -35,13 +35,22 @@ env:
   container-registry-gb200: 766267172432.dkr.ecr.us-east-2.amazonaws.com
 
 jobs:
+  ephemeral-runner-routing:
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_ephemeral_runner_routing.yml@d0878a142525ff79366ca96fbc0b306cd8d64670 # v1.9.0
+    with:
+      lane: ${{ vars.NEMO_CI_EPHEMERAL_LANE }}
+      legacy_label: nvidia-ci-aws-gpu-x8-ephemeral
+      run_id: ${{ github.run_id }}
+      run_attempt: ${{ github.run_attempt }}
+
   is-not-external-contributor:
+    needs: [ephemeral-runner-routing]
     runs-on: ubuntu-latest
     if: github.repository == 'NVIDIA/Megatron-LM'
     outputs:
       is_external_contributor: ${{ github.event.pull_request.user.type == 'User' }}
       is_maintainer: ${{ steps.check-membership.outputs.is_maintainer }}
-      selected_runner: ${{ steps.check-membership.outputs.is_maintainer == 'true' && 'nvidia-ci-aws-gpu-x8' || 'nvidia-ci-aws-gpu-x8-ephemeral' }}
+      selected_runner: ${{ steps.check-membership.outputs.is_maintainer == 'true' && 'nvidia-ci-aws-gpu-x8' || needs.ephemeral-runner-routing.outputs.runner_label }}
       selected_runner_gb200: ${{ steps.check-membership.outputs.is_maintainer == 'true' && 'nvidia-ci-aws-use2-gpu-x4' || 'ubuntu-latest' }}
     permissions:
       issues: write
@@ -168,7 +177,7 @@ jobs:
   pre-flight:
     needs: [is-not-external-contributor]
     if: github.repository == 'NVIDIA/Megatron-LM'
-    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@c1a0837f362a1a696e647238ab1cf916b2a7cf4a # v1.8.6
+    uses: NVIDIA-NeMo/FW-CI-templates/.github/workflows/_cicd_preflight.yml@d0878a142525ff79366ca96fbc0b306cd8d64670 # v1.9.0
 
   configure:
     runs-on: ubuntu-latest
@@ -1357,26 +1366,6 @@ jobs:
               body: `🔄 Merge queue validation started!\n\nYou can track the progress here: ${runUrl}`
             });
 
-  cleanup-taint-node:
-    runs-on: ${{ needs.is-not-external-contributor.outputs.selected_runner }}
-    needs:
-      - is-not-external-contributor
-      - cicd-container-build
-      - cicd-unit-tests-latest
-      - cicd-integration-tests-latest-h100
-      - cicd-integration-tests-latest-gb200
-      - Coverage
-      - Coverage_Fake
-    if: |
-      always()
-      && !cancelled()
-      && contains(needs.is-not-external-contributor.outputs.selected_runner, 'ephemeral')
-      && !needs.pre-flight.outputs.is_deployment_workflow == 'true'
-    steps:
-      - name: Taint node for cleanup
-        shell: bash
-        run: taint-node.sh
-
   DCO_merge_group:
     name: DCO
     if: github.event_name == 'merge_group'

examples/inference/README.md

@@ -111,7 +111,7 @@ Inference can collect traces two ways.
 
 | Path | Enable with | Captures | CUDA graphs |
 |------|-------------|----------|------------|
-| **Sink** | `--moe-enable-routing-replay` | top-K indices only | on |
+| **Sink** | `--moe-enable-routing-replay` + legacy scheduling | top-K indices only | on |
 | **Hook** | no replay + `--cuda-graph-impl none` | indices **+ hidden states + router weights** | must be off |
 
 Only the hook path captures the hidden states and router weights that
@@ -134,10 +134,14 @@ Forward hooks do not fire during CUDA graph replay. MoE cudagraphs must be disab
 
 **Inference — sink** (routing indices only, graphs on):
 
+Routing replay requires legacy scheduling. Because async scheduling is enabled by
+default, explicitly select legacy mode when enabling the routing replay sink.
+
 ```bash
 --moe-routing-trace-path /path/to/trace_dir
 --moe-routing-trace-max-inference-steps 200
 --moe-enable-routing-replay
+--inference-dynamic-batching-async-sched-mode legacy
 ```
 
 **Inference — hook** (adds hidden states + weights for predictability):

examples/inference/advanced/gpt_dynamic_inference.py

@@ -506,7 +506,7 @@ def escape_str(s):
         p_count = len(p_times)
         d_count = len(d_times)
 
-        p_mean = p_total / p_count
+        p_mean = p_total / p_count if p_count != 0 else 0.0
         d_mean = d_total / d_count if d_count != 0 else 0.0
 
         # Commented out for now as the step/add/output times are not calculated correctly.

examples/post_training/modelopt/README.md

@@ -20,7 +20,7 @@ knowledge distillation, pruning, speculative decoding, and more.
 
 ## Major Features
 
-- Start from Hugging Face pretrained model checkpoint with on-the-fly conversion to Megatron-LM checkpoint format.
+- Optimize Megatron-Core distributed checkpoints.
 - Support all kinds of model parallelism (TP, EP, ETP, PP).
 - Export to TensorRT-LLM, vLLM, and SGLang ready unified checkpoint.
 
@@ -50,13 +50,58 @@ Alternatively, you can install from [source](https://github.com/NVIDIA/Model-Opt
 to try our latest features.
 
 > **❗ IMPORTANT:** The first positional argument (e.g. `meta-llama/Llama-3.2-1B-Instruct`) of each script
-> is the config name used to match the supported model config in `conf/`. The pretrained HF checkpoint should
-> be downloaded and provided through `${HF_MODEL_CKPT}`.
+> is the config name used to match the supported model config in `conf/`.
+
+### Megatron-Core Checkpoint Prerequisite
+
+All examples in this folder, including examples for post-training quantization (PTQ) and quantization-aware training (QAT) or distillation (QAD), require a Megatron-Core
+distributed checkpoint. Direct Hugging Face import in the ModelOpt scripts is deprecated; use the
+[Megatron-Bridge stable checkpoint conversion CLI](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/conversion#2-stable-checkpoint-conversion-cli)
+to import the model first. If you want to start directly from a Hugging Face checkpoint for
+quantization, use the [Megatron-Bridge quantization examples](https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization).
+
+The [NVIDIA NGC NeMo container catalog](https://catalog.ngc.nvidia.com/orgs/nvidia/-/containers/nemo/-/tags)
+lists the current release. The example below uses `26.08`; the NeMo container provides
+Megatron-Bridge at `/opt/Megatron-Bridge`. It mounts a local Hugging Face model and a directory
+for the imported checkpoint.
+
+```sh
+HF_MODEL_PATH=/path/to/hf/Llama-3.2-1B-Instruct
+MEGATRON_MODEL_PATH=/path/to/checkpoints/Llama-3.2-1B-Instruct
+HF_TOKEN=<your_hugging_face_access_token>
+
+docker run --rm -it --shm-size=24g \
+    -e HF_TOKEN \
+    -v "${HF_MODEL_PATH}:/models/hf/Llama-3.2-1B-Instruct:ro" \
+    -v "${MEGATRON_MODEL_PATH}:/models/megatron/Llama-3.2-1B-Instruct" \
+    -w /opt/Megatron-Bridge \
+    nvcr.io/nvidia/nemo:26.08 \
+    ./scripts/conversion/convert.sh import \
+        --executor local --device cpu \
+        --hf-model /models/hf/Llama-3.2-1B-Instruct \
+        --megatron-path /models/megatron/Llama-3.2-1B-Instruct
+```
+
+#### Choosing CPU or GPU for the import
+
+Use `--device cpu` for a single-process import when the model and conversion working set fit in
+host memory. This is the simplest choice for the Llama 3.2 1B example above and does not require
+GPUs or model parallelism. Use `--device gpu` with `--gpus-per-node` (and, when needed, `--nodes`)
+when host memory is insufficient, when the target checkpoint needs tensor, pipeline, or expert
+parallelism, or when a distributed import is more practical for a large model. Model size alone
+does not require GPU import: a sufficiently high-memory CPU node can import even a large model;
+GPU import shards the work and memory across ranks. For GPU import, set `--tp`, `--pp`, and `--ep`
+to the intended checkpoint layout.
+
+Set `HF_TOKEN` to a Hugging Face access token before starting the import; it is required for gated
+models and is passed into the container by `-e HF_TOKEN`. Set `--tp`, `--pp`, and `--ep` for the
+target parallelism. Consult the Megatron-Bridge conversion documentation for supported architectures
+and distributed examples.
 
 
 ### ⭐ NVFP4 Quantization, Quantization-Aware Training, and Model Export
 
-Provide the pretrained checkpoint path through variable `${HF_MODEL_CKPT}` and provide variable
+Provide the imported Megatron-Core checkpoint through `${MLM_MODEL_CKPT}` and provide variable
 `${MLM_MODEL_SAVE}` which stores a resumeable Megatron-LM distributed checkpoint. To export
 Hugging Face-Like quantized checkpoint for TensorRT-LLM, vLLM, or SGLang deployement,
 provide `${EXPORT_DIR}` to `export.sh`.
@@ -71,13 +116,12 @@ provide `${EXPORT_DIR}` to `export.sh`.
 ```sh
 \
     TP=1 \
-    HF_MODEL_CKPT=<pretrained_model_name_or_path> \
+    MLM_MODEL_CKPT=<megatron_core_checkpoint> \
     MLM_MODEL_SAVE=/tmp/Llama-3.2-1B-Instruct_quant \
     ./quantize.sh meta-llama/Llama-3.2-1B-Instruct NVFP4_DEFAULT_CFG 
 
 \
     PP=1 \
-    HF_MODEL_CKPT=<pretrained_model_name_or_path> \
     MLM_MODEL_CKPT=/tmp/Llama-3.2-1B-Instruct_quant \
     EXPORT_DIR=/tmp/Llama-3.2-1B-Instruct_export \
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
@@ -89,7 +133,6 @@ export the model with flag `--export-vllm-fq`:
 \
     PP=1 \
     MLM_EXTRA_ARGS=--export-vllm-fq \
-    HF_MODEL_CKPT=<pretrained_model_name_or_path> \
     MLM_MODEL_CKPT=/tmp/Llama-3.2-1B-Instruct_quant \
     EXPORT_DIR=/tmp/Llama-3.2-1B-Instruct_export \
     ./export.sh meta-llama/Llama-3.2-1B-Instruct
@@ -115,14 +158,13 @@ via the auto-quantize arguments.
 ```sh
 \
     TP=1 \
-    HF_MODEL_CKPT=<pretrained_model_name_or_path> \
+    MLM_MODEL_CKPT=<megatron_core_checkpoint> \
     MLM_MODEL_SAVE=/tmp/Llama-3.2-1B-Instruct_auto_quant \
     MLM_EXTRA_ARGS="--auto-quantize-bits 4.0" \
     ./quantize.sh meta-llama/Llama-3.2-1B-Instruct auto
 
 \
     PP=1 \
-    HF_MODEL_CKPT=<pretrained_model_name_or_path> \
     MLM_MODEL_CKPT=/tmp/Llama-3.2-1B-Instruct_auto_quant \
     EXPORT_DIR=/tmp/Llama-3.2-1B-Instruct_auto_quant_export \
     ./export.sh meta-llama/Llama-3.2-1B-Instruct

examples/post_training/modelopt/convert.sh

@@ -27,6 +27,9 @@ if [ -z ${MLM_MODEL_CKPT} ]; then
     if [ -z ${HF_MODEL_CKPT} ]; then
         HF_MODEL_CKPT=${1}
     fi
+    printf "${MLM_WARNING} Importing a Hugging Face checkpoint from convert.sh is deprecated.\n"
+    printf "Use Megatron-Bridge to create a Megatron-Core checkpoint, then set ${PURPLE}MLM_MODEL_CKPT${WHITE}.\n"
+    printf "For Hugging Face quantization workflows, see https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization\n"
     ${LAUNCH_SCRIPT} ${SCRIPT_DIR}/convert_model.py \
         ${MODEL_ARGS} \
         --tensor-model-parallel-size ${TP} \

examples/post_training/modelopt/convert_model.py

@@ -38,7 +38,10 @@ def add_convert_args(parser):
     """Add additional arguments for ModelOpt checkpoint convertion."""
     group = parser.add_argument_group(title='ModelOpt MCore checkpoint convertion')
     group.add_argument(
-        "--pretrained-model-path", type=str, default=None, help="HuggingFace pretrained model"
+        "--pretrained-model-path",
+        type=str,
+        default=None,
+        help="Deprecated: Hugging Face pretrained model. Use Megatron-Bridge to import checkpoints.",
     )
     group.add_argument(
         "--extra-model-path", type=str, default=None, help="Extra module weights to load"
@@ -132,6 +135,14 @@ def check_arguments():
     unwrapped_model = unwrap_model(model)[0]
 
     if args.pretrained_model_path is not None:
+        warnings.warn(
+            "Importing Hugging Face checkpoints with --pretrained-model-path is deprecated. "
+            "Use Megatron-Bridge to create a Megatron-Core checkpoint, then load it with --load. "
+            "For Hugging Face quantization workflows, see "
+            "https://github.com/NVIDIA-NeMo/Megatron-Bridge/tree/main/examples/quantization.",
+            FutureWarning,
+            stacklevel=2,
+        )
         import_dtype = torch.float16 if args.fp16 else torch.bfloat16
         unwrapped_model = unwrap_model(model)[0]
         workspace_dir = os.environ.get("MLM_WORK_DIR", "/tmp")

examples/rl/benchmark_refit.py

@@ -75,12 +75,12 @@ def model_provider(pre_process=True, post_process=True, parallel_output=False,
     )
 
 
-def create_refit_service(method):
+def create_refit_service(method, execution_batch_bytes: int | None = None):
     """Create and return a refit service instance."""
     if method == 'nvshmem':
         return NVSHMEMCopyService()
     elif method == 'nccl_m2n':
-        return NCCLM2NCopyService()
+        return NCCLM2NCopyService(max_group_bytes=execution_batch_bytes)
     elif method == 'nccl':
         return NCCLCopyService()
     elif method == 'gloo':
@@ -104,14 +104,26 @@ def print_config_summary(args, src_config, dst_config, world_size, mode):
     print_rank_0(f"{'='*80}\n")
 
 
-def run_benchmark(src_model, dst_model, refit_service, num_warmup, num_iterations):
+def run_benchmark(
+    src_model,
+    dst_model,
+    refit_service,
+    num_warmup,
+    num_iterations,
+    execution_batch_bytes: int | None = None,
+):
     """Run warmup and benchmark iterations, return timings."""
     # Warmup (builds refit plan on first iteration)
     print_rank_0(f"Warmup: {num_warmup} iterations...")
     for i in range(num_warmup):
         torch.cuda.synchronize()
         torch.distributed.barrier()
-        swap_model_weights(src_model, dst_model, refit_method=refit_service)
+        swap_model_weights(
+            src_model,
+            dst_model,
+            refit_method=refit_service,
+            execution_batch_bytes=execution_batch_bytes,
+        )
         torch.cuda.synchronize()
         torch.distributed.barrier()
 
@@ -126,7 +138,12 @@ def run_benchmark(src_model, dst_model, refit_service, num_warmup, num_iteration
         torch.distributed.barrier()
 
         start_time = time.perf_counter()
-        swap_model_weights(src_model, dst_model, refit_method=refit_service)
+        swap_model_weights(
+            src_model,
+            dst_model,
+            refit_method=refit_service,
+            execution_batch_bytes=execution_batch_bytes,
+        )
         torch.cuda.synchronize()
         end_time = time.perf_counter()
 
@@ -251,7 +268,9 @@ def benchmark_collocated():
 
     # Create refit service
     print_rank_0(f"Creating {args.refit_method} service...")
-    refit_service = create_refit_service(args.refit_method)
+    refit_service = create_refit_service(
+        args.refit_method, execution_batch_bytes=args.refit_execution_batch_bytes
+    )
     print_rank_0("Service created.\n")
 
     try:
@@ -261,6 +280,7 @@ def benchmark_collocated():
             refit_service,
             args.num_benchmark_warmup,
             args.num_benchmark_iterations,
+            execution_batch_bytes=args.refit_execution_batch_bytes,
         )
         print_results(timings)
     finally:
@@ -361,7 +381,9 @@ def benchmark_non_collocated():
 
     # Create refit service
     print_rank_0(f"Creating {args.refit_method} service...")
-    refit_service = create_refit_service(args.refit_method)
+    refit_service = create_refit_service(
+        args.refit_method, execution_batch_bytes=args.refit_execution_batch_bytes
+    )
     print_rank_0("Service created.\n")
 
     try:
@@ -371,6 +393,7 @@ def benchmark_non_collocated():
             refit_service,
             args.num_benchmark_warmup,
             args.num_benchmark_iterations,
+            execution_batch_bytes=args.refit_execution_batch_bytes,
         )
         print_results(timings)
     finally:

megatron/core/inference/config.py

@@ -556,8 +556,9 @@ class InferenceConfig:
     enabled), then all DP ranks share the same sampling / generation seed.
     """
 
-    async_sched_mode: AsyncScheduleMode = AsyncScheduleMode.LEGACY
-    """Mode used to schedule dynamic batching inference work."""
+    async_sched_mode: AsyncScheduleMode = AsyncScheduleMode.ASYNC
+    """Mode used to schedule dynamic batching inference work. Defaults to async scheduling; use
+    ``AsyncScheduleMode.LEGACY`` to disable it."""
 
     logprobs_mode: Literal['raw_logprobs', 'processed_logprobs'] = 'raw_logprobs'
     """Whether returned log-probs are modified by the sampling parameters or not."""