Merge branch 'main' into Stop-logging-checkpoints-as-wandb-artifacts-and-keep

Author: shurkat-nvidia

.github/copy-pr-bot.yaml

@@ -1,4 +1,4 @@
 enabled: true
 auto_sync_draft: false
 auto_sync_ready: true
-trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "XiaodaNV", "YangFei1990", "ZhiyuLi-Nvidia", "abhinav-khattar", "adistomar", "ahmadki", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "jupiterepoch", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwangjunahn", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "nathan-nvidia", "niyunsheng", "ntajbakhsh", "nvcsathe", "nvegesna-netizen", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "wanyingw", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"]
+trustees_override: ["AAnoosheh", "ArEsKay3", "Autumn1998", "BestJuly", "BoxiangW", "CarlosGomes98", "ChenhanYu", "Connor-XY", "DanialTaheri", "FDecaYed", "HaochenYuan", "ISEEKYAN", "JF-D", "JRD971000", "Leili", "Mellonta", "Phlip79", "QiZhangNV", "RPrenger", "ShriyaRishab", "WanZzzzzz", "Wohox", "XiaodaNV", "YangFei1990", "YigongQin", "ZhiyuLi-Nvidia", "abhinav-khattar", "adistomar", "ahmadki", "alokpathy", "ananthsub", "anlthms", "aroshanghias-nvd", "ashehper", "asolergi-nv", "athitten", "balasaajay", "buptzyb", "chtruong814", "cjld", "cspades", "cuichenx", "deepakn94", "desh2608", "dimapihtar", "dingqingy-nv", "duncanriach", "ehosseiniasl", "erhoo82", "ericharper", "fanshiqing", "faradawn", "fitsumreda", "freewym", "frsun-nvda", "gautham-kollu", "gdengk", "goelarushi", "guihong-nv", "guyueh1", "hexinw-nvidia", "huvunvidia", "hxbai", "ilml", "jalbericiola", "janEbert", "jaredcasper", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "jupiterepoch", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwangjunahn", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "nathan-nvidia", "niyunsheng", "ntajbakhsh", "nvcsathe", "nvegesna-netizen", "parthmannan", "philipcmonk", "prajwal1210", "pthombre", "rapatel", "rhewett-nv", "rogerwaleffe", "sajadn", "sancha", "santhnm2", "santiagoakle", "sbak5", "shanmugamr1992", "sharathts", "sheliang-nv", "shengf-nv", "shifangx", "shjwudp", "sidsingh-nvidia", "skyw", "sraman-rgb", "sudhakarsingh27", "svcnvidia-nemo-ci", "tdene", "theothermike", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "wanyingw", "wdykas", "wujingyue", "xiaoyao0115", "xuantengh", "xuwchen", "yaox12", "yaoyu-33", "yashaswikarnati", "yeyu-nvidia", "yobibyte", "youngeunkwon0405", "yqwangustc", "yueshen2016", "yuzhongw-nvidia", "zhehuaichen", "zhongbozhu"]

.github/oncall_schedule.json

@@ -1,8 +1,4 @@
 [
-    {
-        "user": "guihong-nv",
-        "date": "2026-10-07"
-    },
     {
         "user": "ilml",
         "date": "2026-10-14"
@@ -46,5 +42,9 @@
     {
         "user": "guihong-nv",
         "date": "2026-12-23"
+    },
+    {
+        "user": "ilml",
+        "date": "2026-12-30"
     }
 ]

.github/workflows/_build_ci_container.yml

@@ -143,7 +143,8 @@ jobs:
       fail-fast: false
       matrix: ${{ fromJSON(needs.cicd-compute-build-matrix.outputs.matrix) }}
     runs-on: >-
-      ${{ inputs.is-maintainer && matrix.cloud == 'gb-gpu' && 'nvidia-ci-aws-cpu-arm64'
+      ${{ inputs.is-maintainer && matrix.cloud == 'aws-h100' && 'nvidia-ci-aws-cpu'
+          || inputs.is-maintainer && matrix.cloud == 'gb-gpu' && 'nvidia-ci-aws-cpu-arm64'
           || matrix.runner }}
     timeout-minutes: ${{ inputs.timeout-minutes }}
     env:

docker/Dockerfile.ci.dev

@@ -88,13 +88,43 @@ RUN --mount=type=bind,source=docker/common/install_nccl.sh,target=/opt/install_n
     esac
 EOF
 
+ARG CUBLAS_VERSION=13.8.1.7
+RUN bash -euxo pipefail <<"EOF"
+    # cuBLAS 13.8.1 includes the grouped GEMM correctness fix.
+    case "${TARGETARCH}" in
+        amd64)
+            CUBLAS_ARCH=x86_64
+            ;;
+        arm64)
+            CUBLAS_ARCH=sbsa
+            ;;
+        *) echo "Unsupported TARGETARCH for cuBLAS: ${TARGETARCH}" >&2; exit 1 ;;
+    esac
+    CUBLAS_ARCHIVE="libcublas-linux-${CUBLAS_ARCH}-${CUBLAS_VERSION}_cuda13-archive"
+    CUBLAS_TMP=$(mktemp -d)
+    curl -fLSs --retry 3 -o "${CUBLAS_TMP}/cublas.tar.xz" \
+        "https://developer.download.nvidia.com/compute/cublas/redist/libcublas/linux-${CUBLAS_ARCH}/${CUBLAS_ARCHIVE}.tar.xz"
+    tar -xJf "${CUBLAS_TMP}/cublas.tar.xz" -C "${CUBLAS_TMP}"
+    # Replace the toolkit copies so existing linker paths pick up both cuBLAS and cuBLASLt.
+    CUDA_ROOT="${CUDA_HOME:-/usr/local/cuda}"
+    cp -a "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/include/." "${CUDA_ROOT}/include/"
+    cp -a "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/lib/." "${CUDA_ROOT}/lib64/"
+    install -Dm644 "${CUBLAS_TMP}/${CUBLAS_ARCHIVE}/LICENSE" "${CUDA_ROOT}/doc/cublas/LICENSE"
+    ldconfig
+    rm -rf "${CUBLAS_TMP}"
+EOF
+
 COPY pyproject.toml uv.lock /workspace/
 COPY docker/patches/mamba-cxx20.patch /workspace/mamba-cxx20.patch
 COPY megatron/core/__init__.py /workspace/megatron/core/
 COPY megatron/core/package_info.py /workspace/megatron/core/
 ARG IMAGE_TYPE=dev
 ENV IMAGE_TYPE=${IMAGE_TYPE}
 ENV NVTE_BUILD_NUM_PHILOX_ROUNDS=3
+# cuDNN Frontend is replaced after the sync below with this ref of NVIDIA/cudnn-frontend
+# (feat/1.29.0-megatron). Docker-only: outside pyproject.toml and uv.lock. Set empty to keep the
+# locked PyPI wheel. Skipped for the LTS image.
+ARG CUDNN_FRONTEND_GIT_REF=8924d4b41fd678ba1b98cd7f98bb78b222637573
 # Preserve Ampere support unless a caller selects narrower TE build targets.
 ARG NVTE_CUDA_ARCHS="80;90;100"
 # TE compiles each source's GPU targets one after another by default (nvcc --threads 1), so
@@ -183,6 +213,10 @@ RUN --mount=type=cache,id=uv-${TARGETARCH}-${NVTE_CUDA_ARCHS},target=/root/.cach
     rm -rf "${MAMBA_BUILD_DIR}"
     trap - EXIT
 
+    if [ -n "$CUDNN_FRONTEND_GIT_REF" ] && [ "$IMAGE_TYPE" != "lts" ]; then
+        uv pip install --no-deps --reinstall \
+            "nvidia-cudnn-frontend @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_GIT_REF}"
+    fi
 EOF
 
 # Reuse the wheel without copying a second archive into the final image.

docs/api-guide/core/generalized_tensor_parallel.md

@@ -49,6 +49,7 @@ Both `GTP_remat` collectives are prefetched one step ahead, so they overlap the
     - [2.5 Tuning knobs](#25-tuning-knobs)
     - [2.6 FP32-accumulation wgrad reduce-scatter (optional)](#26-fp32-accumulation-wgrad-reduce-scatter-optional)
     - [2.7 NCCL symmetric-memory wgrad reduce-scatter (optional)](#27-nccl-symmetric-memory-wgrad-reduce-scatter-optional)
+    - [2.8 Batch and sequence partitioning within GTP_remat](#28-batch-and-sequence-partitioning-within-gtp_remat)
   - [3. Implementation details](#3-implementation-details)
     - [3.1 GTP\_remat architecture (Mcore ↔ TE integration)](#31-gtp_remat-architecture-mcore--te-integration)
       - [What the flags do under the hood](#what-the-flags-do-under-the-hood)
@@ -228,7 +229,10 @@ See [§3.3 Distributed checkpointing (DCP)](#33-distributed-checkpointing-dcp) f
 
 ## 2. Usage
 
-GTP_remat is enabled through two CLI flags on Megatron's training launcher; everything else (process-group construction, parameter slicing, prefetch chain wiring, optimizer routing) is automatic once the flags are set.
+Dense and expert weight-shard flags configure GTP_remat independently. An optional
+sequence-shard flag controls sequence partitioning within dense weight groups.
+Process-group construction, parameter slicing, prefetch chain wiring, and optimizer
+routing are handled automatically from these settings.
 
 ### 2.1 Knob summary
 
@@ -238,6 +242,7 @@ The table below covers every GTP-related CLI flag and Python knob. "Required" me
 |---|---|---|---|---|
 | `--tensor-parallel-num-weight-shards` | **Required** | Always, to activate dense GTP | — | Total TP×GTP_remat shards per dense weight; GTP_remat degree = value ÷ TP. Must be ≥ TP and divisible by it. [§2.2](#22-required-flags) |
 | `--expert-tensor-parallel-num-weight-shards` | **Required** | MoE models (to shard routed-expert weights) | — | Total ETP×EGTP_remat shards per expert weight; EGTP_remat degree = value ÷ ETP. Independent of dense axis. [§2.2](#22-required-flags) |
+| `--tensor-parallel-num-sequence-shards` | **Optional** | Split sequences within dense GTP groups | TP with SP, otherwise 1 | Includes TP/SP. Within each GTP_remat group, ranks are partitioned into batch-shard groups. Each group processes an independent microbatch, with its ranks handling different sequence partitions. [§2.8](#28-batch-and-sequence-partitioning-within-gtp_remat) |
 | `--gtp-remat-reduce-scatter-with-fp32-accumulation` | **Optional** | BF16 wgrads **and** GTP_remat axis ≥ 4 | off | Replaces the ring RS with an all-to-all + local FP32 sum to eliminate per-hop rounding error. Auto-bypassed at axis size ≤ 2. [§2.6](#26-fp32-accumulation-wgrad-reduce-scatter-optional) |
 | `--gtp-remat-nccl-ub` | **Optional** | For enabling symmetric-memory NCCL kernels on supported systems | off | Enables symmetric memory registration for the dense gtp_remat wgrad reduce-scatter path. Takes precedence over fp32-accum on its group; incompatible with `--disable-symmetric-registration`. [§2.7](#27-nccl-symmetric-memory-wgrad-reduce-scatter-optional) |
 | `--gtp-expert-remat-nccl-ub` | **Optional** | For enabling symmetric-memory NCCL kernels on supported systems | off | Enables symmetric memory registration for the routed-expert egtp_remat wgrad reduce-scatter path. [§2.7](#27-nccl-symmetric-memory-wgrad-reduce-scatter-optional) |
@@ -432,6 +437,35 @@ it** — a different collective over a different process group, so enable either
   end-to-end at 256 GPUs; the VMM allocator recovers it). Window registration accepts this
   memory and runs the same symmetric kernels.
 
+### 2.8 Batch and sequence partitioning within GTP_remat
+
+`--tensor-parallel-num-sequence-shards S` splits sequences within the dense weight
+group. Enabling GTP sequence sharding with TP > 1 requires `--sequence-parallel`;
+`S` includes TP and must be a multiple of TP and divide the weight-shard count.
+Leaving it unset preserves existing TP/SP behavior. `--context-parallel-size`
+adds sequence partitions across weight replicas.
+
+For example, on 128 GPUs with PP=1:
+
+```bash
+--tensor-model-parallel-size 1 \
+--tensor-parallel-num-weight-shards 64 \
+--tensor-parallel-num-sequence-shards 64 \
+--context-parallel-size 2
+```
+
+This gives DP=1, with weights sharded across GTP_remat groups of 64 ranks and each
+sequence partitioned across 128 ranks. The batch dimension is unpartitioned
+(`batch_parallel_size=1`), so all 128 ranks process different sequence partitions
+of the same microbatch.
+
+This example assumes each NVLink domain can accommodate a 64-rank weight-sharding
+group. Use rank ordering with TP as the innermost axis and PP as the outermost axis,
+and place each weight-sharding group entirely within one NVLink domain.
+
+GTP sequence sharding requires Megatron DDP and `--ckpt-format torch_dist`;
+dynamic CP, FSDP, and sequence-packing schedulers are unsupported.
+
 ---
 
 ## 3. Implementation details
@@ -458,9 +492,29 @@ it** — a different collective over a different process group, so enable either
 
 The `--*-num-weight-shards` flags flow through five stages, from process groups to the prefetch chain:
 
-1. **Process groups.** `initialize_model_parallel(...)` treats GTP_remat/EGTP_remat as **first-class orthogonal axes** (`world = TP·CP·GTP_remat·DP`; experts `= ETP·EP·EGTP_remat·PP·expert_dp`), building `_GTP_WEIGHT_REMAT_GROUP` and `_EXPERT_GTP_WEIGHT_REMAT_GROUP` (sizes = `num-weight-shards / TP` and `/ ETP`). CP is placed more locally (smaller stride) than GTP_remat on the dense/decoder axis, since CP's collectives are on the model's critical path (`_inject_gtp_remat_axis(..., after="cp")`). **DP and gtp_remat stay orthogonal:** `get_data_parallel_group()` is the replicate axis (DDP + optimizer shard over it); `with_gtp_remat=True` gives the combined DP × gtp_remat axis for data distribution.
+1. **Process groups.** With GTP sequence sharding, `GTPParallelLayout` gives the same ranks two views:
 
-   > **Batch-size arithmetic.** `args.data_parallel_size` is the **replicate degree only** — gtp_remat is *divided out* of it (folded into `total_model_size` at `arguments.py:446`). But data is distributed over the **full DP × gtp_remat axis**, so each gtp_remat peer consumes a *distinct* microbatch and the global sample count is `micro_batch_size × data_parallel_size × gtp_weight_remat_size × num_microbatches`. The training loop therefore **re-applies `gtp_weight_remat_size`** to close the gap: *multiplied back in* for the LR-scheduler `increment` and the logged `batch_size`, *divided back out* to recover `eval_num_microbatches`. Without this it would read as a double-count — it is not.
+   - **Weights:** TP × GTP_remat shards each weight; configured CP × DP replicates those shards.
+   - **Tokens:** configured CP × `gtp_remat_num_sequence_shards` splits sequences; the remaining GTP_remat ranks × DP process different samples.
+
+   `gtp_remat_num_sequence_shards` excludes TP/SP. Validation multiplies CLI CP by this
+   count to obtain runtime `context_parallel_size`; direct Python callers supply
+   that combined size. Bootstrap constructs the groups from these views; without
+   GTP sequence sharding, it uses the existing independent-axis layout. The expert
+   grid remains ETP × EP × EGTP_remat × PP × expert-DP.
+
+   The GTP_remat axis splits into sequence and batch partitions:
+
+   ```text
+   GTP_remat_size = gtp_remat_num_sequence_shards × gtp_remat_num_batch_shards
+   batch_parallel_size = DP × gtp_remat_num_batch_shards
+   num_microbatches = global_batch_size / (micro_batch_size × batch_parallel_size)
+   ```
+
+   Here `gtp_remat_num_batch_shards` is the derived batch extent, not a separate
+   flag. By default, `gtp_remat_num_sequence_shards=1`, so all GTP_remat ranks
+   distribute independent samples. Global and micro batch sizes come from the CLI;
+   the microbatch calculator uses `batch_parallel_size` to derive `num_microbatches`.
 
 2. **Per-class sharding.** `extensions/transformer_engine.py` decides *per linear class* whether to shard, so **no `gtp_remat_group` is threaded through the module APIs** (attention, Mamba, MLP, embedding, MTP). Dense wrappers resolve the group via `utils.get_gtp_weight_remat_group(...)`; `TEGroupedLinear` uses `pg_collection.expt_gtp_remat`. Group `None`/size-1 → left full; otherwise `_gtp_pre_init` pre-shards `out_features` and `_gtp_attach_post_init` makes the shard a **`GTPShardedParam`** (the `DistributedWeight` implementer; native FP8/NVFP4 by reclass, BF16 by re-register). Base `te.Linear` (MoE latent projections) receives a group only when `--gtp-remat-opt-in-modules moe_latent_proj` is set; otherwise it stays full → see [Class hierarchy](#class-hierarchy-which-linears-shard).
 
@@ -539,7 +593,9 @@ Under **full-iteration CUDA graphs** the recompute-forward is captured; `wait_as
 
 ![DDP + (E)GTP_remat interaction with the distributed optimizer](../../images/generalized_tensor_parallel/0611_ddp_egtp_orthogonal_bucketing.png)
 
-**(E)GTP_remat is *super loosely coupled* to DDP and the distributed optimizer — they stay almost completely GTP_remat-agnostic.** GTP_remat is just another sub-axis of the rank grid (`world = TP×CP×GTP_remat×DP`); a GTP_remat-sharded weight rides the *exact same* code path as an ordinary param. There are **no** GTP_remat/EGTP_remat-specific buffers, optimizers, or bucket groups, and just **one** GTP_remat-specific gradient-scaling factor (the expert-buffer correction below). The entire DDP/DistOpt stack touches GTP_remat in only **four** narrow places:
+GTP_remat weights use the existing DDP buffers and optimizer paths. DDP reduces over
+weight replicas, while GTP_remat completes reductions across weight shards. This also
+applies when those ranks process sequence partitions. The integration has four parts:
 
 1. **finalize all-reduce** (`_allreduce_replicated_grads_over_gtp_remat_group`) — completes the gtp_remat axis for *replicated* (non-GTP_remat) params (SUM under `calculate_per_token_loss`, AVG otherwise; see §3.2 table); a no-op when GTP_remat is inactive.
 2. **`is_gtp_weight_remat` / `allreduce` tags** propagated onto the optimizer's master shards — consumed only by the grad-norm dedup filter.
@@ -886,12 +942,23 @@ Case A is what §1.3's "tail slice" framing describes for the reassembled tensor
 torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parallel/ -v -m "not flaky_in_dev"
 ```
 
+Run the sequence-sharding tests on both 4 and 8 GPUs to cover all layouts:
+
+```bash
+for ranks in 4 8; do
+    uv run python -m torch.distributed.run --nproc-per-node "$ranks" -m pytest -q \
+        tests/unit_tests/test_gtp_parallel_layout.py \
+        tests/unit_tests/generalized_tensor_parallel/test_gtp_cp_overlap.py
+done
+```
+
 | Test file | What it guards |
 |-----------|----------------|
 | `test_gtp_basics.py` | Core GTP_remat shard/gather, cache ownership, wgrad ring, DDP bucket alignment, the `num_zeros` padding correction (§3.7), and TE's recompute-phase flag (dtype-agnostic, relied on by the readiness gate in §3.2). |
 | `test_attention_gtp.py` | GTP_remat on attention linears, loss parity vs no-GTP_remat. |
 | `test_mamba_gtp.py` | GTP_remat on Mamba projection weights. |
 | `test_tp_gtp.py` | GTP_remat composed with tensor parallelism (`tp_group × gtp_remat_group`). |
+| `test_gtp_cp_overlap.py` | Sequence sharding with CP, TP/SP, gradient and Adam-update parity, and checkpoint loading after a CP change. |
 | `test_moe_egtp.py` | EGTP_remat on MoE routed-expert weights. |
 | `test_gtp_loss_correctness.py` | End-to-end: GTP_remat per-step loss trajectory matches a no-GTP_remat baseline. |
 | `test_gtp_grad_correctness.py` | Gradient + dist-opt + grad-norm numeric parity vs a DP baseline at replicate (DP) > 1. Also the fp32-accumulation reduce-scatter (§2.6): gtp_remat-axis and DDP-axis parity, plus the size-2 bypass. |