[PyTorch] Propagate skip_fp8_weight_update in GroupedLinear during FP8 CUDA graph capture

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch

[LeSingh1]

Description

GroupedLinear.forward hardcoded None for skip_fp8_weight_update, so the FP8 graph-capture skip tensor was never forwarded during CUDA graph replay — the bug reported in #3051. Linear.forward handles this correctly: when FP8GlobalStateManager.fp8_graph_capturing() is true it reads quantization_state.skip_fp8_weight_update_tensor, forces is_first_microbatch = False, and threads the tensor into its forward args. GroupedLinear omitted that retrieval block and passed a literal None, even though the autograd _GroupedLinear.forward already unpacks and uses the flag (quantize_weight(..., skip_update_flag=skip_fp8_weight_update)).

Fix

Mirror the Linear reference path in GroupedLinear.forward: add the same fp8_graph_capturing() retrieval block and thread skip_fp8_weight_update into non_tensor_args (positionally aligned with the existing autograd unpack). FP8GlobalStateManager is already imported in the module.

Verification status

Runtime-unverified — developed without a CUDA GPU, so the FP8 / CUDA-graph numerics and the Nemotron-MoE gradient-suppression reproducer were not run. The change is a near-mechanical parity fix against the established Linear path (same accessor, same is_first_microbatch forcing, same threading), and the positional unpack lines up with _GroupedLinear.forward. A GPU regression test (comparing GroupedLinear FP8 grads under graph-replay-active vs delayed) belongs alongside the existing CUDA-graph tests but can't run on this platform — I'd ask a maintainer with a GPU to confirm numerics before merge, and I'm happy to add the test.

Developed with AI assistance.

Addresses #3051

[greptile-apps]

Greptile Summary

This PR fixes the missing propagation of skip_fp8_weight_update in GroupedLinear.forward during FP8 CUDA graph capture, mirroring the pattern already used in Linear.forward. A dedicated regression test is added to verify that graphed and non-graphed FP8 runs produce identical outputs with weight caching enabled.

• Core fix (grouped_linear.py): Inserts an fp8_graph_capturing() retrieval block before m_splits validation to set skip_fp8_weight_update and force is_first_microbatch = False during CUDA graph replay — the same pattern used by Linear.
• New test (test_cuda_graphs.py): Adds _GroupedLinearWrapper to adapt the 3-D test harness, a new grouped_linear branch in _test_cuda_graphs, and a dedicated parametrized test that asserts graphed outputs match eager outputs under both DelayedScaling and Float8CurrentScaling recipes.

Confidence Score: 5/5

Safe to merge after resolving the duplicate fp8_graph_capturing() block; the fix itself is correct and the new test exercises the right code path.

The fix correctly mirrors Linear.forward and threads skip_fp8_weight_update into non_tensor_args, verified by the new regression test. The only issue is a redundant duplicate of the same retrieval block that was independently introduced into main before the PR was merged — both blocks compute identical values, so there is no incorrect behavior. The test logic is sound for the given model config.

transformer_engine/pytorch/module/grouped_linear.py — contains the duplicate retrieval block at lines 1687–1694 and 1708–1715 that should be reduced to one.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/module/grouped_linear.py	Adds fp8_graph_capturing() retrieval block before m_splits processing; identical block already existed after m_splits in the base branch, leaving both in the file as redundant duplicates
tests/pytorch/test_cuda_graphs.py	Adds _GroupedLinearWrapper, grouped_linear branch in _test_cuda_graphs, and a focused parametrized regression test; model config produces 64 total tokens (divisible by num_gemms=2) so the assertion holds

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GroupedLinear.forward called] --> B{fp8_graph_capturing?}
B -- Yes --> C[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
B -- No --> D[skip_fp8_weight_update = None]
C --> E{skip_fp8_weight_update
is not None?}
D --> E
E -- Yes --> F[is_first_microbatch = False]
E -- No --> G[is_first_microbatch unchanged]
F --> H[Validate m_splits]
G --> H
H --> I{fp8_graph_capturing?
⚠️ Duplicate block}
I -- Yes --> J[skip_fp8_weight_update =
quantization_state.skip_fp8_weight_update_tensor]
I -- No --> K[skip_fp8_weight_update = None]
J --> L{skip_fp8_weight_update
is not None?}
K --> L
L -- Yes --> M[is_first_microbatch = False again]
L -- No --> N[is_first_microbatch unchanged]
M --> O[Build non_tensor_args with skip_fp8_weight_update]
N --> O
O --> P[_GroupedLinear.apply / forward]
style I fill:#ffcccc
style J fill:#ffcccc
style K fill:#ffcccc
style L fill:#ffcccc
style M fill:#ffcccc

Loading

_{Reviews (6): Last reviewed commit: "Merge branch 'main' into fix-grouped-lin..." | Re-trigger Greptile}

[ptrendx]

@LeSingh1 Could you add the test that would exercise this functionality?

[LeSingh1]

@ptrendx Added a test in tests/pytorch/test_cuda_graphs.py: test_make_graphed_callables_grouped_linear_with_fp8_weight_caching.

It wraps GroupedLinear (flattening the [seqlen, batch, hidden] harness input to 2D and splitting evenly across 2 GEMMs) so it slots into the existing _test_cuda_graphs harness, then runs the full / individual / none graph modes with fp8_weight_caching=True and asserts the graphed outputs equal the eager reference. That equality only holds when is_first_microbatch is threaded into the weight-update skip tensor on every microbatch — on the pre-fix None hardcode the cached FP8 weights diverge, so the test fails without the change. Parametrized over DelayedScaling + Float8CurrentScaling, both fp8_params settings, and the fp32/fp16 dtypes already used in this file.

One disclosure: I don't have FP8-capable GPU hardware locally, so I verified the test by construction (mirroring test_make_graphed_callables_with_fp8_weight_caching) rather than by running it — it'll need a CI/GPU run to confirm. Happy to adjust the recipe coverage or wrapper approach if you'd prefer it structured differently.

[LeSingh1]

Rebased onto main to resolve a conflict with #3038 (the GroupedLinear graph-safe refactor). Worth noting: #3038 plumbs skip_fp8_weight_update all the way through _GroupedLinear's autograd function, but the module-level GroupedLinear.forward still passes None, # skip_fp8_weight_update into non_tensor_args, so the value never reaches that plumbing during graph capture. This change wires it up the same way Linear.forward does (read the skip tensor from FP8GlobalStateManager.quantization_state when fp8_graph_capturing()). The test from the previous commit comes along with the rebase. Still GPU/CI-only on my end since I don't have FP8 hardware locally.

[ptrendx]

/te-ci pytorch

[ksivaman]

/te-ci pytorch