Pass the owner's collection to FP8/FP4 quantization contexts by Connor-XY · Pull Request #8035 · NVIDIA/Megatron-LM

Connor-XY · 2026-10-09T21:39:48Z

Depends on #8034, which adds amax_reduction_group and the pg_collection keyword. Until it merges, please review the four commits after #8034's, one per owner, in the comparison with #8034's branch. Both branches merged main, so that comparison's file view also shows #8034's changes.

Why. These components hold their model's collection, but their FP8/FP4 autocasts still reduce amaxes over the global group, or over WORLD when there is no global grid:

transformer blocks and layers, MTP, recompute;
HybridStack;
CUDA graph runners;
the multimodal projector and the TE wrappers.

On 4 GPUs, a layer whose own TP×DP×CP ranks are [0, 2] reduces over [0, 1] on main.

What. Each owner passes pg_collection=:

TE wrappers, the LM head and the projector. TELinear gains an optional pg_collection.
Block, layer, recompute and MTP. MultiTokenPredictionBlock and T5Model add tp_cp and tp_dp_cp to their default collections.
HybridStack. MambaLayer keeps its collection.
The CUDA graph runner. TECudaGraphHelper uses amax_reduction_group, which removes a dead branch that read a nonexistent tp_dp field.

Tests.

Per-owner tests record the group of each autocast, with the global accessor patched to raise. They fail on main.
In the integration run (all [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 PRs merged onto main 77391e7, against main), the four new per-owner test files pass on 4 GPUs: test_te_amax_reduction_group.py (7 tests), test_quantization_context_amax_group.py (12), test_hybrid_stack_amax_group.py (3) and test_cuda_graph_amax_group.py (2). Take the FP8/FP4 amax reduction group from a process-group collection #8034's test_amax_reduction_group.py, which this branch carries, passes 23 tests there, one of them added by another [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 PR.
In the integration run, this PR's test files show no regressions against main. The touched existing files pass 14 tests vs 5 on main (test_schedule_quantization_context.py), 5 vs 4 (test_multimodal_projector.py) and 67 vs 67 (test_hybrid_block.py) on 4 GPUs, and 31 vs 13 (test_t5_model.py) on 8 GPUs (two nodes). The extra passes include other [TRACKING] Deprecate parallel_state: explicit process groups in megatron/core #6307 PRs' new tests.
Before the 2026-11-01 sync, FP8/FP4 training (delayed, tensorwise, MXFP8, NVFP4, TE CUDA graphs) gave the same results as main.
Before the sync, the FP8, hybrid and TE regression sets gave the same results as main.

Compatibility. On the standard grid the collection holds the same group objects, so numerics are unchanged. A caller-built collection without tp_dp_cp (or without tp_cp when tp_only_amax_red is set) now raises ValueError under FP8/FP4. Needs determinism-exempt.

Synced with main 77391e7 on 2026-11-01 by a merge commit, so existing commits and review links stay valid. The merge had no conflicts; its message lists the test and allowlist updates main required.

Part of #6307.

🤖 Generated with Claude Code

Connor-XY · 2026-10-09T21:39:50Z

/ok to test 1dd9473

copy-pr-bot · 2026-10-09T21:39:52Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

get_fp8_context, get_fp4_context and the Transformer Engine per-module quantization autocast read the amax reduction group from parallel_state whenever model parallelism is initialized, and pass no group otherwise. A model built on its own process-group collection therefore reduces amaxes over another model's TP x DP x CP ranks or, without the global grid, over Transformer Engine's default group. Add amax_reduction_group(pg_collection, tp_only_amax_red), the collection counterpart of get_amax_reduction_group(with_context_parallel=True). It returns tp_cp when tp_only_amax_red is set and tp_dp_cp otherwise, and raises ValueError when that field is not set. tp_dp_cp spans the gtp_remat axis like the global group, and on the standard grid both fields hold the same group objects that parallel_state returns. get_fp8_context, get_fp4_context and the per-module quantization autocast take an optional keyword-only pg_collection and reduce over the group it selects. The per-module path picks the field with each recipe's own tp_only_amax_red, so training and evaluation recipes may differ. Callers that pass no collection keep the global group. The group is resolved only for autocast contexts, because fp8_model_init does not reduce amaxes. black 26 reformats one untouched raise in get_fp4_recipe; it is included so that the changed file passes the format check. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Regenerate the allowlist from the rebased tree. The base's file predates three upstream changes to the reads in finalize_model_grads.py, cuda_graphs.py and dsa_gqa.py; this series adds or removes none. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Transformer Engine wrappers that apply a per-module quantization recipe, the MXFP8 LM head and the multimodal projector are built with a process-group collection, but their FP8/FP4 autocasts still reduce amaxes over the global group, or over no group without the global grid. Pass the collection to the per-module quantization autocast and to get_fp8_context: - TELinear takes an optional pg_collection. TEColumnParallelLinear, TERowParallelLinear, the LM head and InferenceLinear forward theirs. - TELayerNormColumnParallelLinear and TEGroupedLinear keep theirs. - MultimodalProjector passes its collection, or None when it has none, which keeps the global group. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

TransformerBlock, TransformerLayer, the activation-recompute path and the MTP layers hold a process-group collection, but open FP8/FP4 autocasts that reduce amaxes over the global group. A model on its own grid then mixes its amaxes with another model's TP x DP x CP ranks, or reduces over no group when the global grid does not exist. Pass the owner's collection to get_fp8_context and get_fp4_context from the block's outer and per-layer contexts, the recompute path, the layer's fine-grained context and the MTP contexts. The MTP layer keeps its collection for this. MultiTokenPredictionBlock and T5Model build default collections from the global grid with only the groups they read themselves. Add tp_cp and tp_dp_cp to them so FP8/FP4 runs that pass no collection keep reducing over the same groups. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

HybridStack opens FP8/FP4 autocasts around the stack and around each layer that reduce amaxes over the global group, although the stack holds the model's process-group collection. Pass the collection to get_fp8_context and get_fp4_context. MambaLayer keeps its collection as pg_collection, like TransformerLayer, so code that receives a graphable layer can read the layer's groups. The hybrid block tests build a collection with only tp, pp and cp; add the amax reduction groups, which their FP8 cases now need. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

A local CUDA graph runner re-enters the layer's FP8/FP4 autocast while it captures and warms up, and reduces amaxes over the global group. Pass the graphed layer's collection instead. TECudaGraphHelper already took the group from its collection through a private helper whose with_context_parallel=False branch read a tp_dp field that ProcessGroupCollection does not have; only the other branch was ever called. Use amax_reduction_group, which raises ValueError instead of asserting when the field is missing. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY · 2026-10-26T22:29:20Z

/ok to test be10fd5

The combined 1F1B delayed FP8 context passes the model chunk's collection to get_fp8_context, whose pg_collection keyword comes from NVIDIA#8034, which NVIDIA#8035 contains. The new tests sit on NVIDIA#8035's edit of test_schedule_quantization_context.py and use the amax-group test helpers that NVIDIA#8034 adds. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8028 and NVIDIA#8035 each refreshed the allowlist against main. Merging them keeps NVIDIA#8035's total of 195, while the merged entries record 194 reads: NVIDIA#8028 removes the read in forward_step_calc_loss that NVIDIA#8035 still counted. Regenerate the file so that its total matches its entries. No entry changes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8035 takes the FP8/FP4 amax reduction group from a process-group collection: TELinear accepts pg_collection= and the TE linear wrappers keep it as _pg_collection. NVIDIA#8748 passes the owners' collections to the TE linears they build and reads fused_layer._pg_collection when it splits a TELayerNormColumnParallelLinear, so without NVIDIA#8035 these constructions raise TypeError or AttributeError under Transformer Engine. Merge NVIDIA#8035's head (which contains NVIDIA#8034). The merge is clean, including the allowlist. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The merges of NVIDIA#8017 and NVIDIA#8035 took NVIDIA#8017's optimizer entries (one get_megatron_optimizer shim read, from NVIDIA#6259) and NVIDIA#8035's total of 195, while the merged entries record 194 reads. Regenerate the file from the tree; only the total changes, and the result equals NVIDIA#8017's allowlist. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

NVIDIA#8035 gives each layer the FP8/FP4 amax reduction group of its collection, through NVIDIA#8034's amax_reduction_group, and passes it to the CUDA graph runners. The FP8/FP4 contexts and the CUDA-graph helpers are to warn when they fall back to the global amax group, in the code NVIDIA#8034 and NVIDIA#8035 change, so merge NVIDIA#8035 here. Four files conflict with NVIDIA#6293 and NVIDIA#6303, which this branch carries: - transformer_engine.py: keep the imports of both sides. - t5_model.py: keep this branch's side. T5Model already resolves an omitted collection through LanguageModule (NVIDIA#6303), which includes the amax groups NVIDIA#8035 adds. - multi_token_prediction.py: keep the fallback warning and the hcp group, and add NVIDIA#8035's tp_cp and tp_dp_cp to required_pgs. - test_hybrid_block.py: get_pg_collection keeps the union of both sides, NVIDIA#6303's embd group and NVIDIA#8035's tp_cp and tp_dp_cp amax reduction groups, with NVIDIA#8035's comment. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded. NVIDIA#8755 warns the global fallbacks of the log and param-hash helpers, RoPE, MIMO partitioning, batch-invariant unpermute, the FP8/FP4 contexts and the CUDA-graph helpers, and brings NVIDIA#8017, NVIDIA#8032, NVIDIA#8035, NVIDIA#8047, NVIDIA#8049 and NVIDIA#8053 (NVIDIA#6303 already came in with NVIDIA#8752). megatron/core/extensions/transformer_engine.py conflicts on its parallel_state import: NVIDIA#8038 drops get_tensor_model_parallel_world_size and keeps model_parallel_is_initialized, and NVIDIA#8755's eac0b1c drops get_amax_reduction_group and model_parallel_is_initialized. The merged file calls none of the three, so the import keeps only get_context_parallel_group, get_hierarchical_context_parallel_groups and get_tensor_model_parallel_group. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The checker classification that follows rewrites NVIDIA#8766's deprecated function kind and NVIDIA#8765's accessor deprecation, which NVIDIA#8767 brings, counts NVIDIA#8747's registry getter as a shim, and edits the compatibility-page lines and the accessor-use test lines that NVIDIA#8767 adds. NVIDIA#8767 brings NVIDIA#7029, which this branch does not have yet. Three conflicts are resolved as follows: - megatron/core/extensions/transformer_engine.py, parallel_state import: keep NVIDIA#7029's TeardownStage and register_model_parallel_teardown, and drop get_amax_reduction_group, get_tensor_model_parallel_world_size and model_parallel_is_initialized, which NVIDIA#8755 and NVIDIA#8038 stop calling. - The same file, TEColumnParallelLinear and TERowParallelLinear: keep NVIDIA#8753's _gtp_replica_group(...) for gtp_replica_group (NVIDIA#8767 brings NVIDIA#8753 through NVIDIA#8761), plus NVIDIA#8035's pg_collection=pg_collection argument that NVIDIA#8755 brings. - megatron/core/transformer/cuda_graphs.py: keep NVIDIA#7029's parallel_state import and NVIDIA#8755's amax_reduction_group and resolve_process_groups import; the ProcessGroupCollection import that NVIDIA#8755's resolve_process_groups call replaces is dropped. The process-group usage allowlist is regenerated on the merged tree, with the _comment that NVIDIA#8766's checker, which NVIDIA#8767 brings, writes with --update. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

The checker classification that follows requires every allowlisted read in megatron/core to be warned or guarded. NVIDIA#8769 warns before resolve_gtp_remat_group reads the global grid, and brings NVIDIA#8734 (optimizer, duplicate-filter and tensor-parallel seeding fallback warnings), NVIDIA#8057, NVIDIA#8060 and NVIDIA#8073. Five files conflict. Each is resolved as follows: - optimizer/optimizer.py, process_groups_config import: take NVIDIA#8752's require_process_groups and resolve_process_groups (92d2030) and NVIDIA#8734's _FALLBACK_WARNING_SKIP_PREFIXES, ProcessGroupFallbackWarning, resolve_gtp_remat_group and warn_global_process_group_fallback (ec801ee). - generalized_tensor_parallelism.py: keep NVIDIA#7029's TeardownStage and register_model_parallel_teardown import, which NVIDIA#8767 brings, and NVIDIA#8734's resolve_process_groups import. - inference_layers.py: InferenceLinear keeps NVIDIA#8035's comment and pg_collection=pg_collection, which NVIDIA#8755 brings, and passes NVIDIA#8734's owner=type(self).__name__ to resolve_gtp_remat_group. - tensor_parallel/layers.py: NVIDIA#8734 stops calling the four parallel_state rank and size getters, so their import goes; NVIDIA#8734's process_groups_config import stays, and get_global_memory_buffer comes from megatron.core.utils as NVIDIA#8730, which NVIDIA#8767 brings, has it. - tests/unit_tests/ssm/test_hybrid_block.py get_pg_collection: the required groups are the union of NVIDIA#8035's embd, tp_cp and tp_dp_cp and NVIDIA#8073's ep, expt_tp and tp_ep, with both comments. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…8748 The fallbacks that follow give the Attention, MTP and MoE owners the GTP groups of their linears, and the one-warning tests extend NVIDIA#8748's test_owner_linear_pg_collection.py: a Transformer Engine owner without a collection must warn once, which needs NVIDIA#8748 to pass the owner's collection to its linears. NVIDIA#8748 brings NVIDIA#8017 and NVIDIA#8035. NVIDIA#8035, which NVIDIA#8748 brings, meets NVIDIA#6293, which NVIDIA#8729 brings, in three files. They are resolved as follows: - transformer_engine.py: the process_groups_config import takes both sides (amax_reduction_group and warn_global_process_group_fallback). - t5_model.py: keep NVIDIA#6293's side. T5Model already resolves an omitted collection with every global group at the top of its constructor, which includes the amax groups NVIDIA#8035 adds. - multi_token_prediction.py: keep NVIDIA#6293's fallback warning and hcp group, and add NVIDIA#8035's tp_cp and tp_dp_cp to required_pgs. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

…8739 This base holds every other PR of the parallel_state deprecation series (NVIDIA#6307): the end-to-end test that follows asserts that no process-group fallback or deprecation warning fires anywhere on a collection that is not the global grid, and the strict-fallback functional tests turn every remaining fallback into an error, so both need the whole series below them. NVIDIA#8739 brings NVIDIA#8028 and NVIDIA#8035. NVIDIA#8035, which NVIDIA#8739 brings, meets NVIDIA#6293, which NVIDIA#8736 brings through NVIDIA#8058, in three files. They are resolved as follows: - transformer_engine.py: the process_groups_config import takes both sides (amax_reduction_group and warn_global_process_group_fallback). - t5_model.py: keep NVIDIA#6293's side. T5Model already resolves an omitted collection with every global group at the top of its constructor, which includes the amax groups NVIDIA#8035 adds. - multi_token_prediction.py: keep NVIDIA#6293's fallback warning and hcp group, and add NVIDIA#8035's tp_cp and tp_dp_cp to required_pgs. The process-group usage allowlist is regenerated on the merged tree. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Adapted to main outside the conflicts: - tests/unit_tests/transformer/test_quantization_context_amax_group.py: The bare MultiTokenPredictionLayer in test_projection_and_layer_contexts_use_the_layer_collection now sets is_hybrid_mtp = False instead of mtp_layer_pattern = None. Main's NVIDIA#7555 removed mtp_layer_pattern, and _proj_and_transformer_layer now branches on is_hybrid_mtp, which only __init__ sets. The GPT branch matches the test's mtp_model_layer stub. The process-group allowlist records the reads main added since this PR was written. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yan Xu <yxu1@nvidia.com>

Connor-XY · 2026-11-03T14:45:24Z

Superseded by #8893, which carries this change as commit 976a872 under the same title, with "Originally #8035." in its message. Closing so review happens in one place. The branch stays on the fork.

Connor-XY added the determinism-exempt label Oct 9, 2026

copy-pr-bot Bot temporarily deployed to test October 9, 2026 21:41 Inactive

copy-pr-bot Bot temporarily deployed to public October 9, 2026 21:44 Inactive

lhb8125 mentioned this pull request Oct 12, 2026

Megatron-LM Monthly Activity Report — October 2026 (Final, through October 31, 2026) #7882

Open

Connor-XY and others added 6 commits October 26, 2026 01:30

Connor-XY force-pushed the yx/pstate-b2-amax-callers branch from 1dd9473 to be10fd5 Compare October 26, 2026 22:25

copy-pr-bot Bot temporarily deployed to test October 26, 2026 22:30 Inactive

copy-pr-bot Bot temporarily deployed to public October 26, 2026 22:35 Inactive

Connor-XY mentioned this pull request Oct 30, 2026

Warn on the last unwarned core reads and once per omitted collection #8772

Closed

This was referenced Nov 1, 2026

Pass owners' process-group collections to the linears they build #8748

Closed

Warn on optimizer, TP and helper fallbacks to WORLD or the global grid #8734

Closed

Connor-XY closed this Nov 3, 2026

Navigation Menu

Repository navigation

Pass the owner's collection to FP8/FP4 quantization contexts - #8035

Pass the owner's collection to FP8/FP4 quantization contexts#8035
Connor-XY wants to merge 7 commits into
NVIDIA:mainfrom
Connor-XY:yx/pstate-b2-amax-callers

Connor-XY commented Oct 9, 2026 •

edited

Loading

Uh oh!

Connor-XY commented Oct 9, 2026

Uh oh!

copy-pr-bot Bot commented Oct 9, 2026

Uh oh!

Connor-XY commented Oct 26, 2026

Uh oh!

Connor-XY commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

Connor-XY commented Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Connor-XY commented Oct 9, 2026

Uh oh!

copy-pr-bot Bot commented Oct 9, 2026

Uh oh!

Connor-XY commented Oct 26, 2026

Uh oh!

Connor-XY commented Nov 3, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Connor-XY commented Oct 9, 2026 •

edited

Loading