Graph Safe support for TE Grouped linear Op

[vthumbe1503]

Description

Graph Safe Support for GroupedLinear Op for MXFP8/BF16/FP16:
Group Quantize + CUBLAS GGEMM for the above data types

Context Saving Refactoring for Graph Safe Paths
GroupedLinear and its Fused MOE Variants(GemmSwiglu based MLP) both have unified way of saving context and and hence using something like fused forward and unfused backward works without hassle.

Use save_for_backward for activations(GroupedTensorStorage) as well
GroupedTensorStorage now has prepare_for_saving and restore_from_saved functions defined
And this is used to save activations for both unfused GroupedLinear module and GEMM+Swiglu based MLP for mxfp8

Megatron Main Grad logic refactoring
All the repeated logic of using Megatrong DDP/FSDP specific attributes from parameters are abstracted out into reusable functions.

Fixes # (issue)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR adds graph-safe (CUDA-graph-compatible) support for the GroupedLinear op by introducing a new _fuser_forward_grouped_tensor / _fuser_backward_grouped_tensor code path that dispatches to general_grouped_gemm_for_grouped_tensor, gated on Blackwell (SM100+) hardware, supporting MXFP8, BF16, and FP16. It also refactors the Megatron-LM main-grad logic into shared helpers and aligns the fused MLP saved-tensor layout so the unfused backward can consume the same context.

• New graph-safe forward/backward paths in grouped_linear.py gate on _is_graph_safe_path_supported (SM100+) and dispatch to a GroupedTensor-based GEMM; the legacy split_quantize path is preserved as a fallback.
• Megatron main-grad helpers are extracted into _common.py and adopted across basic_linear.py, backward_grouped_mlp.py, backward_linear_add.py, backward_linear_scale.py, and userbuffers_backward_linear.py.
• GroupedTensorStorage gains prepare_for_saving / restore_from_saved helpers, and the fused MLP forward/backward layouts are unified with the standalone GroupedLinear layout.

Confidence Score: 5/5

The change is safe to merge. The new graph-safe forward/backward paths are well-structured, delay_wgrad=True correctly returns None for weight gradients and defers the GEMM to backward_dw, and the Megatron main-grad refactor is a clean consolidation with no functional regression.

The core new code paths correctly handle all major branches: quantized vs. unquantized, single vs. per-group weights, delay_wgrad, and accumulate_into_main_grad. The shared _common.py helpers are straightforward extractions of duplicated inline logic. The unified saved-tensor layout for the fused MLP forward/backward is internally consistent and round-trips correctly through the unfused backward path.

No files require special attention for merge safety. grouped_linear.py has pre-existing issues carried over from earlier reviews that the team is tracking separately.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/ops/basic/grouped_linear.py	Core change: adds graph-safe grouped-tensor forward/backward, new helper methods, and fuser_forward_save_ctx; forward-pass main_grad guard at lines 894-898 is still a debug-level placeholder (flagged in earlier reviews); main_grad.view(-1) non-contiguous risk in _fuser_backward_grouped_tensor line 1624 also previously flagged. New code paths look logically correct including delay_wgrad=True returning None and wgrad GEMM executed via wgrad_store.pop().
transformer_engine/pytorch/ops/_common.py	Adds four utility helpers for Megatron main-grad handling; logic is correct and cleanly consolidated from duplicated inline code across several backward ops.
transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py	Saved-tensor layout updated to [split_sizes, base_split_offsets, split_points, grouped_x, *weights]; old per-component unpacking replaced with direct GroupedTensor save. References to .data corrected to .rowwise_data.
transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py	Saved-tensor layout and ctx attribute names unified with the standalone GroupedLinear grouped-tensor forward; use_grouped_tensor_path=True set so the unfused backward can consume the same ctx.
transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py	Adds prepare_for_saving / restore_from_saved helpers; new API is defined but the primary save/restore paths in this PR use direct field-nulling instead, creating two divergent patterns for future contributors.
transformer_engine/pytorch/ops/basic/basic_linear.py	Main-grad helpers refactored to shared functions; logic is equivalent to the removed inline code.
transformer_engine/pytorch/ops/fused/backward_linear_add.py	Main-grad inline logic replaced with shared helpers; no functional change.
transformer_engine/pytorch/ops/fused/backward_linear_scale.py	Same refactor as backward_linear_add.py; no functional change.
transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py	Main-grad inline logic replaced with shared helpers; error label updated from BasicLinear to UserbuffersBackwardLinear for clarity.

_{Reviews (10): Last reviewed commit: "fix linting errors" | Re-trigger Greptile}

[greptile-apps]

Missing contiguity error handling for main_grad.view(-1)

main_grad.view(-1) will raise a generic RuntimeError if main_grad is non-contiguous (e.g. when returned by get_main_grad() via __fsdp_param__). The equivalent code in backward_grouped_mlp.py wraps the reshape in try/except and re-raises with an actionable message that includes the shape and stride. Without that guard, users hitting this case will see an opaque PyTorch error instead of a clear diagnostic.

[greptile-apps]

Non-contiguous main_grad causes opaque RuntimeError in graph-safe backward

main_grad.view(-1) will raise a plain RuntimeError whenever main_grad is non-contiguous (e.g., when FSDP returns a main_grad that is already shaped (num_groups, out_features, in_features) but lives as a non-unit-stride slice of a larger gradient buffer). view_main_grad_as_grouped_buffer only guards the reshape-to-grouped-shape step — once the helper returns the tensor unchanged (shape already matches), the subsequent .view(-1) is unprotected.

The equivalent code in backward_grouped_mlp.py avoids this by passing main_grad directly to make_grouped_tensor_from_rowwise_data, which internally calls .contiguous() when needed. grouped_linear.py should do the same:

Suggested change

	grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data(
	num_tensors=num_groups,
	tensor_shape=weight_shape,
	rowwise_data=main_grad.view(-1),
	dtype=main_grad.dtype,
	)
	grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data(
	num_tensors=num_groups,
	tensor_shape=weight_shape,
	rowwise_data=main_grad,
	dtype=main_grad.dtype,
	)

[timmoon10]

This is an important PR and I think this looks good overall. My main thought is that we need to maintain the contract that unfused ops and fused ops are interchangeable.

[vthumbe1503]

/te-ci pytorch

[timmoon10]

This function is similar to the context setup stage in #2877. After this PR is merged, we should try to consolidate them together.

After that, it may be worthwhile calling this function from the fused ops so that we guarantee the unfused and fused ops have the same context format.