Repository navigation
Conversation
| def _apply_linear_o_group_proj(self, input_: torch.Tensor) -> torch.Tensor: | ||
| """Apply TE BatchedGEMM when available, otherwise use the legacy einsum.""" | ||
| op = self._get_linear_o_group_proj_op() | ||
| if op is None: |
There was a problem hiding this comment.
Can we add a linear_o_group_proj to DSv4HybridSelfAttentionSubmodules spec to select the TE op?
I think the silent fallback makes it not easy to locate some issues.
|
/claude strict-review |
Strict Review SummaryFindings by severity: CRITICAL: 0 · IMPORTANT: 0 · SUGGESTION: 1 AssessmentThis is a well-scoped, well-tested change that swaps the grouped output projection ( I traced the areas most likely to hide bugs and found them handled correctly:
Most impactful noteOne minor [SUGGESTION Simplification] (inline): the einsum-fallback Risk level: Low. No correctness, performance, or backward-compatibility concerns. The change is confined to an experimental attention variant, preserves checkpoint format across both implementations, and is covered by targeted unit tests. |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Xin Yao <xiny@nvidia.com>
7f8f133 to
37937e6
Compare
|
/claude strict-review |
|
Strict Review Summary — PR #5722: Add batched gemm for DSv4HybridAttention Findings by severity — CRITICAL: 1, IMPORTANT: 0, SUGGESTION: 0 This is a carefully engineered PR. Swapping the raw grouped nn.Parameter + torch.einsum for te.pytorch.BatchedLinear (with a clean einsum fallback + warning when TE is too old) is sound, the MXFP8/BF16 numerics + gradient tests are thorough, and the substantial state-dict shim work correctly preserves the legacy 2D checkpoint layout across state_dict/load_state_dict/sharded_state_dict. The forward-path change (apply_linear_o_group_proj) is behavior-preserving, and the einsum equation is unchanged. No backward-compat issue for existing checkpoints in the plain state_dict path, and no new direct parallel_state.get*_group() reads. Most impactful finding: [CRITICAL Correctness] The sharded_state_dict build_fn unconditionally reshapes to the full 2D checkpoint_local_shape. Under the DEFAULT distributed-optimizer sharding (fully_sharded_model_space), the optimizer wraps this factory around flat, partial DP-shard state with flattened_range set — a full-shape reshape either raises a numel-mismatch RuntimeError (weight straddling a DP-shard boundary) or produces 2D data with a stale non-None flattened_range. The established precedent (ssm/mamba_mixer.py::_split_tensor_factory) handles this by narrowing/keeping data flat and passing flattened_range through. Suggested fix: skip the reshape when flattened_range is not None (the two leading axes merge contiguously, so flat element order is identical), and add a distopt save/load test — the current test only exercises factory.build() with full-shape data. Risk assessment: Medium. The core BatchedLinear integration and non-distributed checkpoint paths are well-tested and low-risk. The risk concentrates in the one distributed-optimizer save path above, which is the default distopt sharding type and is not covered by the added tests. Recommend resolving the CRITICAL finding (and adding distopt coverage) before merge. |
| def build_fn( | ||
| key: str, tensor: torch.Tensor, replica_id: ReplicaId, flattened_range: Optional[slice] | ||
| ) -> ShardedTensor: | ||
| checkpoint_tensor = tensor.reshape(checkpoint_local_shape) |
There was a problem hiding this comment.
[CRITICAL Correctness] build_fn unconditionally reshapes to the full 2D checkpoint_local_shape, which breaks distributed-optimizer checkpoint save under the default sharding type (fully_sharded_model_space → sharded_param_state_fs_model_space; see distrib_optimizer.py line 1503 default and 1563).
In that path the optimizer wraps this factory around flat, partial state via replace(sharded_metadata, data=state_ten, flattened_range=item_slice) (distrib_optimizer.py lines 2026–2035), where state_ten is a 1D DP-shard slice whose numel is param_range.end - param_range.start — i.e. only the portion of this parameter residing on the local DP rank. When the grouped-proj weight straddles a DP-shard boundary, tensor.reshape(checkpoint_local_shape) raises a RuntimeError (numel mismatch); even when it doesn't straddle, you end up with 2D-shaped data carrying a non-None flattened_range, which is inconsistent for the save machinery.
Contrast:
sharded_param_state_fully_reshardablereshapes to the full shape before wrapping and assertsflattened_range is None(lines 1770, 1805), so a full-shape reshape is valid there.- The established reshape/split factory in
ssm/mamba_mixer.py(_split_tensor_factory) handles the flat case by usingShardedTensor.narrowand passingflattened_rangethrough — it never reshapes flat partial data to the full multi-dim shape.
The previous plain-nn.Parameter implementation produced a normal ShardedTensor that the optimizer handled natively, so this is a regression on the default distopt save path for DSv4HybridAttention.
Fix: the (groups, rank, in)→(groups*rank, in) reshape merges the two leading axes contiguously, so the flat element order is identical. When flattened_range is not None, skip the reshape and keep tensor flat, building the checkpoint ShardedTensor with local_shape=checkpoint_local_shape and the passed-through flattened_range (mirror the mamba/narrow pattern). Only reshape in the non-flattened case. Please also add a distributed-optimizer save/load test (the current sharded_state_dict test only exercises factory.build() with full-shape weight data, not the flat optimizer-shard path).
There was a problem hiding this comment.
The previous implementation did assume full-shaped factory data, so it would fail if invoked with a flat partial optimizer state.
After tracing the checkpoint entry point, I changed the fix slightly instead of adding DSv4-specific support for
fully_sharded_model_space:
fully_sharded_model_spaceis only the fallback whenDistributedOptimizer.sharded_state_dict()is called without checkpoint metadata.- The normal checkpoint path sets
dp_reshardableby default, orfully_reshardablewhen requested.dp_reshardabledoes not invoke this parameter factory, whilefully_reshardablepasses full-shaped state withflattened_range=None.fully_sharded_model_spaceis deprecated and is already incompatible globally becauseShardedTensor.validate_metadata_integrity()rejects every non-Noneflattened_range. Supporting it only for this DSv4 parameter would therefore not make a complete optimizer checkpoint valid.I simplified the implementation accordingly:
- The fallback path now retains the original 2D
Parameter, so it uses the existing state-dict and distributed-checkpoint behavior directly.- Only the TE
BatchedLinearpath installs the 3D-to-2D compatibility factory.- The factory supports the active non-flat checkpoint formats, preserves PP/FSDP group-axis metadata, and explicitly rejects the deprecated flat format.
I also added coverage for the fallback 2D
ShardedTensorpath and pipeline-offset metadata. The full test file passes on GB200 with two ranks (35 passedon each rank), including BF16, MXFP8, quantized model initialization, legacy checkpoint loading, and the einsum fallback.If
fully_sharded_model_spaceneeds to be restored, I think it should be addressed globally in the distributed-checkpoint/optimizer implementation rather than with a DSv4-specific irregular-shard conversion.
Signed-off-by: Xin Yao <xiny@nvidia.com>
|
/claude strict-review |
There was a problem hiding this comment.
Strict review passed — no significant issues found. LGTM
Traced the full change: TE BatchedLinear replaces the hand-rolled Parameter+einsum grouped output projection, with a warning-guarded einsum fallback for older TE.
Verified:
- Backward-compat checkpoints: the legacy 2D
linear_o_group_projkey is preserved on save (state-dict hook) and load (_load_from_state_dictreshape) in both directions; dist-checkpoint mapping viaShardedTensorFactory(3D group-axis → 2D axis-0). Covered by tests for BF16 and MXFP8 quantized weights, plus asharded_offsetscase. - Init parity:
_initialize_grouped_output_projectionappliesinit_methodon aview(-1, in), matching the old 2D init exactly. - Numerics + gradients: TE path matches the einsum reference for BF16 and MXFP8 (fwd + input/weight grads).
_extra_statedrop is safe: delayed-scaling FP8 is excluded and only MXFP8/BF16 are targeted (no persistent amax history), so nothing is lost.- No new global process-group reads; no unused new identifiers (the
tepylint disable is correctly removed now thatte.pytorchis read).
Findings — CRITICAL: 0, IMPORTANT: 0, SUGGESTION: 0. Risk level: low. Comprehensive unit coverage (numerics, gradients, checkpoint round-trip, sharded state dict, fallback path).
|
/ok to test cbbd9ea |
|
/ok to test 03d463c |
|
/ok to test 1b641a8 |
Signed-off-by: Xin Yao <xiny@nvidia.com>
Signed-off-by: Hongxiao Bai <hongxiaob@nvidia.com>
|
/resolve-conflict |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Xin Yao <xiny@nvidia.com>
…tch_dim rename - Add get_dsv4_o_group_proj_spec() and use it in both the config-aware DSv4 attention spec and HybridModel's static hybrid_stack_spec (csa_layer and csa_qk_layernorm_layer), which previously had no linear_o_group_proj and silently fell back to torch.einsum. - hybrid_dsv4_stack_spec: also override csa_qk_layernorm_layer, so 'C' layers with qk_layernorm use the config-aware DSv4 attention. - Pass input_batch_dim to te.ops.BatchedLinear (renamed from batch_dim in TE). - Tests: cover csa_qk_layernorm_layer, move attention to CUDA before BatchedLinear forward, compare MXFP8 weights on the same device, and patch the static spec in the einsum-fallback test. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
What does this PR do?
Use
transformer_engine.pytorch.modules.BatchedLinearforDSv4HybridAttention.linear_o_group_proj, which.contiguous()call after the group projection.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment @NVIDIA/mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.