Add GDP inference unit tests

.github/CODEOWNERS

@@ -1,3 +1,6 @@
+# Files directly in the repository root without a more specific owner.
+/* @NVIDIA/core-nemo
+
 megatron/core/ @NVIDIA/core-adlr @NVIDIA/core-nemo
 megatron/core/tensor_parallel/generalized_tensor_parallelism.py @NVIDIA/core-adlr @NVIDIA/core-nemo @NVIDIA/gtp
 

.github/actions/check-nvidia-sso-membership/action.yml

@@ -1,5 +1,5 @@
 name: 'Check NVIDIA SSO Membership'
-description: 'Check if a GitHub username exists in the NVIDIA SSO users list from github-audits'
+description: 'Check if a GitHub username is in the NVIDIA SSO users list or internal allowlist'
 author: 'NVIDIA'
 
 inputs:
@@ -24,7 +24,7 @@ inputs:
 
 outputs:
   is_member:
-    description: 'Boolean - true if user is in NVIDIA SSO list, false otherwise'
+    description: 'Boolean - true if user is in NVIDIA SSO list or internal allowlist, false otherwise'
     value: ${{ steps.check-membership.outputs.is_member }}
   is_org_member:
     description: 'Boolean - true if user has NVIDIA or NVIDIA-NeMo in org_roles'
@@ -91,6 +91,15 @@ runs:
 
         echo "Checking if $USERNAME is in NVIDIA SSO users list..."
 
+        # Service accounts cannot enroll in SSO but can be explicitly trusted.
+        if [ "$USERNAME" = "svcnemo-autobot" ]; then
+          echo "$USERNAME found in NVIDIA internal user allowlist"
+          echo "is_member=true" >> $GITHUB_OUTPUT
+          echo "is_org_member=false" >> $GITHUB_OUTPUT
+          echo "user_orgs=" >> $GITHUB_OUTPUT
+          exit 0
+        fi
+
         # Check if SSO file is available
         if [ "${{ steps.download-sso.outputs.sso_file_available }}" != "true" ] || [ ! -f "$SSO_FILE" ]; then
           echo "ERROR: $SSO_FILE not available - cannot check membership"

.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", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "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", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "ntajbakhsh", "nvcsathe", "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", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "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", "YangFei1990", "ZhiyuLi-Nvidia", "adistomar", "ahmadki", "aklife97", "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", "jenchen13", "jiaji-huang", "jiemingz", "jingqiny-99", "jkamalu", "jon-barker", "jstjohn", "kajalj22", "kamran-nvidia", "kevalmorabia97", "kevjshih", "kingformatty", "ko3n1g", "ksivaman", "kunlunl", "kvareddy", "kwyss-nvidia", "lauradang", "layalir", "lhb8125", "liding-nv", "lmcafee-nvidia", "maanug-nv", "macandro96", "mathemakitten", "matthieule", "mchrzanowski", "mehraakash", "minitu", "mkhona-nvidia", "nanz-nv", "niyunsheng", "ntajbakhsh", "nvcsathe", "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", "thomasdhc", "tomlifu", "trintamaki", "tylerpoon", "vasunvidia", "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": "janEbert",
-        "date": "2026-08-05"
-    },
     {
         "user": "maanug-nv",
         "date": "2026-08-12"
@@ -46,5 +42,9 @@
     {
         "user": "janEbert",
         "date": "2026-10-21"
+    },
+    {
+        "user": "maanug-nv",
+        "date": "2026-10-28"
     }
 ]

.gitlab/stages/06.triage.yml

@@ -49,11 +49,17 @@ triage:linear_write:
       artifacts: true
   allow_failure: true
   script:
-    - >-
-      cerno-linear write
-      --config "${CERNO_CONFIG}"
-      --plan linear_action_plan.json
-      --output linear_action_plan_post.json
+    # The shared GitHub email lookup reads the existing GH_TOKEN CI variable.
+    - |
+      ONCALL_ASSIGNEE="$(
+        python tests/test_utils/python_scripts/resolve_oncall_assignee.py \
+          --schedule-file .github/oncall_schedule.json
+      )"
+      cerno-linear write \
+        --config "${CERNO_CONFIG}" \
+        --plan linear_action_plan.json \
+        --output linear_action_plan_post.json \
+        --assignee "${ONCALL_ASSIGNEE}"
   artifacts:
     when: always
     paths:

docker/Dockerfile.linting

@@ -24,7 +24,7 @@ RUN --mount=type=secret,id=JET_INDEX_URLS \
 
 # Keep this in the internal-only stage so public CI has no internal service dependency.
 ARG CI_SERVER_URL
-ARG CERNO_COMMIT=5a5fb5360e67f8f09d189871bbc0d768c09c43fa
+ARG CERNO_COMMIT=3219afcbfcdbddd1d55e013d94e54a15986e2b1f
 RUN --mount=type=secret,id=CERNO_TOKEN \
       GIT_CONFIG_COUNT=1 \
       GIT_CONFIG_KEY_0=http.extraHeader \

docs/api-guide/core/generalized_tensor_parallel.md

@@ -417,7 +417,7 @@ The figure visualizes the per-class split from the list above: green = resolves
 
 Two distinct pools with explicit lifecycle rules:
 
-- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`. Same-shape buffers across layers are shared, **except between chain neighbours** — one-step-ahead keeps `prev_w` and the current weight live at once, so `_ensure_distinct_buffer_from_prev` folds a parity bit into the key when the two would collide, at the cost of one extra buffer for the second of the pair. Normally inert (neighbours are different weight roles, hence different shapes); it fires when CG capture leaves two same-shaped weights adjacent — embedding + output_layer alone in the `UNGRAPHED` chain. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay.
+- **`GTPWeightCache`** (AG/RS output buffers) — ticket-based, keyed on `(shape, dtype, fwd, expert_idx, reduce_scatter)`, plus a `("recompute", parity)` suffix for recompute-chain gathers. Same-shape buffers across layers are shared, **except between chain neighbours** — one-step-ahead keeps the predecessor and the current weight live at once, so `_ensure_no_shared_buffer_with` folds a parity bit into the key when the two would collide, at the cost of one extra buffer for the second of the pair. The caller names which chain to guard, because the chains disagree on who a weight's neighbour is: on the fwd chain the check is normally inert (neighbours are different roles, hence different shapes) and fires only when CG capture leaves two same-shaped weights adjacent — embedding + output_layer alone in the `UNGRAPHED` chain — whereas on a recompute chain same-shape adjacency is the norm. Tickets persistent; buffer allocated lazily on first `get()`; addresses stable across iterations for CG replay.
 - **`_wgrad_buf_pool`** (wgrad-GEMM output recycling) — holds the **full, unsharded** wgrad-GEMM output buffer (shape `_unsharded_shape`, dtype `main_grad.dtype` — fp32 when `grad_reduce_in_fp32`, else bf16). The TE backward writes the wgrad into it via `main_grad_func = weight.grad_buffer` (a `DistributedWeight` protocol method backed by `get_wgrad_tensor`; it is a *scratch*, distinct from the sharded `param.main_grad`); the protocol's `finalize_group_grads` (backed by `wgrad_reduce_scatter`) then reduce-scatters it down to the shard and the buffer is returned here. This is a full-weight-shaped fp32/bf16 transient — one of the larger per-weight buffers — and is **precision-independent** (wgrad is always computed in high precision), so it is identical in BF16 vs MXFP8 runs. Buffers are tagged `_from_gtp_wgrad_pool=True` at `_wgrad_pool_get`; `_wgrad_pool_put` no-ops on foreign buffers (fresh allocs from Megatron `layers.py` or aten F.embedding bwd) → caching allocator handles those, so the pool never accumulates untagged buffers.
 
 #### Overlap design summary
@@ -433,7 +433,7 @@ GTP_remat runs up to **three** independent prefetch chains, all following one ru
 |---|------|---------|--------------------|--------------|------|
 | 1 | fwd | weight `i` | `next_w` = i+1 ‖ `GEMM_i` | rowwise (`fwd=True`) | `_prefetch_handle` |
 | 2 | bwd dgrad | weight `i` | `prev_w` = i−1 ‖ `Dgrad_i` | columnwise (`fwd=False`) | `_prefetch_handle` |
-| 3 | bwd recompute | weight `i` | `_recompute_next` = i+1 ‖ `recompute_GEMM_i` | rowwise (`fwd=True`) | `_recompute_prefetch_handle` (separate) |
+| 3 | bwd recompute | weight `i` | `_recompute_next` = i+1 ‖ `recompute_GEMM_i` | rowwise (`fwd=True`) | `_recompute_prefetch_handle` + `_ag_ticket_recompute` (separate) |
 | 1b | fwd (MoE, eager) | expert weight `i` | same role in MoE block i+1 ‖ *whole block i* | rowwise (`fwd=True`) | `_prefetch_handle` |
 
 Row 1b is chain 1 applied to a *homogeneous* chain: routed-expert `fc1`/`fc2` link across consecutive MoE blocks, so the runway is a full block rather than one GEMM (§3.4 *Grouped-expert chains*).
@@ -452,7 +452,7 @@ A future MR will add an opt-in wgrad-before-dgrad schedule on `_Linear` / `_Laye
 
 ##### Recompute-forward prefetch chain  *(GTP_remat + activation recompute)*
 
-When a GTP_remat-sharded module is in `--recompute-modules` (e.g. `shared_experts`), its forward is **re-run during backward** to regenerate activations. That recompute-forward must all-gather each weight **rowwise** again — a *third* gather lifecycle, concurrent with the in-flight **columnwise** dgrad gather of the *same* weight. Since both share one `GTPShardedParam`, the recompute path gets its **own** prefetch slot (`_recompute_prefetch_handle` / `_recompute_ag_event`, reusing the `_ag_ticket_fwd` rowwise buffer) so it never clobbers the dgrad lifecycle's `state` / `_prefetch_handle` / `ag_event`.
+When a GTP_remat-sharded module is in `--recompute-modules` (e.g. `shared_experts`), its forward is **re-run during backward** to regenerate activations. That recompute-forward must all-gather each weight **rowwise** again — a *third* gather lifecycle, concurrent with the in-flight **columnwise** dgrad gather of the *same* weight. Since both share one `GTPShardedParam`, the recompute path gets its **own** prefetch slot (`_recompute_prefetch_handle` / `_recompute_ag_event`) so it never clobbers the dgrad lifecycle's `state` / `_prefetch_handle` / `ag_event`, and its **own** buffer ticket (`_ag_ticket_recompute`) with a parity of its own. Reusing `_ag_ticket_fwd` is unsafe twice over: a fwd prefetch may still be in flight in that buffer, and the fwd parity is decided against `prev_w` — a different neighbour. Without its own parity, consecutive recompute nodes share one buffer and the one-ahead prefetch overwrites the weight still being read: silent wrong activations, then NaN.
 
 The recompute weights form a **separate** linked list (`_recompute_next`), **self-populated** on the first backward from the weights actually re-gathered while `in_fp8_activation_recompute_phase()` is true — membership is *observed*, not configured (no tagging, so it tracks exactly what each checkpointed module re-gathers). Each recompute-forward consume prefetches the next recompute weight, so every gather **except the global-first** overlaps preceding recompute / dgrad / wgrad compute:
 
@@ -606,7 +606,7 @@ Three consequences:
   - one-block-ahead makes block *N* and block *N+1* weights **live at the same time** — same key, two tensors in flight;
   - fix: a chain-position **parity (0,1,0,1…)** is folded into the cache key, so consecutive blocks alternate between **exactly two** buffers (counter cleared by `reset_gtp_state()`);
   - without it the prefetch would **overwrite the weight the running GEMM is still reading** — a silent-correctness bug, not a crash;
-  - the hazard is not exclusive to grouped chains — *any* chain whose neighbours share a key has it. Grouped chains are same-key throughout, so they take the blanket counter; others take the narrower `_ensure_distinct_buffer_from_prev` check, which allocates only where the collision is real (see *Buffer / memory management*).
+  - the hazard is not exclusive to grouped chains — *any* chain whose neighbours share a key has it. Grouped chains are same-key throughout, so they take the blanket counter; others take the narrower `_ensure_no_shared_buffer_with` check, which allocates only where the collision is real (see *Buffer / memory management*).
 - **Eager only** — the optimization disables itself under CUDA-graph capture:
   - `_classify_param_chain` evaluates `graphed = _FULL_ITERATION or ("moe" in cuda_graph_modules)` **before** the split, and returns the plain `GRAPHED` chain when it is true;
   - so with `--cuda-graph-impl full_iteration` **every** param is `GRAPHED` — expert weights included — and they keep the ordinary one-step-ahead prefetch;
@@ -741,6 +741,7 @@ torchrun --nproc-per-node 4 -m pytest tests/unit_tests/generalized_tensor_parall
 | `test_gtp_partial_cg.py` | Four-layer partial-CG loss and eager-vs-replay grad-norm parity with two-slot ring reuse across independently replayed graphs (§3.5). |
 | `test_gtp_dcp.py` | DCP sharding metadata (§3.3): TP×GTP_remat offsets, pad reshard, `replica_id`, native-FP8 save/load. |
 | `test_gtp_muon_dcp.py` | Muon optimizer-state DCP roundtrip (§1.6): `replica_id` fold + native-FP8 backfill matching. |
+| `test_gtp_recompute_chain.py` | Recompute-chain buffers (§3.1): adjacent nodes never share a gather buffer, dense and grouped, plus dgrad/wgrad parity vs no-recompute. |
 | `test_gtp_mtp.py` | GTP_remat + MTP shared weights (§3.5), 14 cases over `mtp_use_repeated_layer` × dense/MoE. Both MTP hazards are silent, so each needs its own guard: the async reduce-scatter path is compared numerically against the sync path on an identical model/sharding/batch, and all-gathers issued are tallied against consumes to catch a consume reading a buffer nothing gathered into. |
 | `test_gtp_fp8_param_gather.py` | Native-FP8 GTP_remat (§1.3): fp8-vs-BF16 loss parity (TP1/TP2, MoE), post-save-spike guard. |
 | `test_gtp_custom_pgs.py` | `pg_collection` plumbing: a custom `gtp_remat` group (permuted ranks, same size) must give the same fwd/bwd results as the MPU groups — catches modules reading `parallel_state` instead of the collection passed to them. |

examples/academic_paper_scripts/detxoify_lm/generate_samples_gpt.py

@@ -70,6 +70,7 @@ def model_provider(pre_process=True, post_process=True) -> GPTModel:
         pre_process=pre_process,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+        logit_dtype=getattr(args, 'logit_dtype', None),
         parallel_output=False,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
         position_embedding_type=args.position_embedding_type,

gpt_builders.py

@@ -91,6 +91,7 @@ def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None, pg_
         pre_process=pre_process,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+        logit_dtype=getattr(args, 'logit_dtype', None),
         parallel_output=True,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
         position_embedding_type=args.position_embedding_type,

hybrid_builders.py

@@ -33,6 +33,7 @@ def hybrid_builder(args, pre_process, post_process, vp_stage=None, config=None,
         pre_process=pre_process,
         post_process=post_process,
         fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,
+        logit_dtype=getattr(args, 'logit_dtype', None),
         parallel_output=True,
         share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,
         position_embedding_type=args.position_embedding_type,