Add grouped HybridStack overlap support

[Connor-XY]

Summary

This MR adds grouped HybridStack support and extends EP-overlap scheduling/checkpoint compatibility to grouped hybrid layers.

Key changes:

• Add bracketed HybridStack group syntax, e.g. [*-], [*E], M[M*]-.
• Build bracketed groups as nested HybridStack instances.
• Reject invalid recursion such as nested brackets inside a group.
• Keep MoE constrained to the last symbol inside a group for overlap scheduling.
• Rename/migrate TransformerSchedulePlan usage to HybridStackSchedulePlan, preserving backward-compatible aliases.
• Add HybridStack fine-grained callables so grouped hybrid layers can participate in combined 1F1B / EP overlap.
• Make grouped HybridStack checkpoints use Transformer-compatible logical layer keys.
• Make HybridModel.sharded_state_dict() match GPT behavior by dropping empty output_layer._extra_state.

Checkpoint Compatibility

Grouped HybridStack checkpoint keys now canonicalize to logical Transformer layer keys. For example, [*-] maps both attention and MLP weights to:

decoder.layers.N.self_attention.*
decoder.layers.N.mlp.*

instead of using separate physical layer ids. final_norm is also sharded as final_layernorm for Transformer compatibility.

Verified both directions:

• Transformer checkpoint -> HybridModel: load succeeded, one train step completed, checkpoint saved, validation/test completed.
• HybridModel checkpoint -> Transformer: load succeeded, one train step completed, checkpoint saved, validation/test completed.

Cross-load status:

STATUS save_final_hybrid 0
STATUS transformer_ckpt_to_hybrid 0
STATUS final_hybrid_ckpt_to_transformer 0

Validation

unit tests:

python -m pytest -q \
  tests/unit_tests/ssm/test_hybrid_layer_allocation.py \
  tests/unit_tests/ssm/test_hybrid_block.py \
  tests/unit_tests/models/test_hybrid_model.py::TestHybridModel::test_grouped_sharded_state_dict_uses_transformer_checkpoint_keys

Result:

86 passed, 63 warnings

Also ran:

• git diff --check
• DeepSeek-V3 deterministic Transformer/Hybrid baseline and EP-overlap smoke tests
• Transformer VPP4 baseline/overlap comparison
• strict model-key checkpoint cross-load checks using --dist-ckpt-strictness raise_unexpected

Notes

The checkpoint cross-load validation intentionally uses model-weight / finetune-style loading with optimizer and RNG state skipped. Full non-finetune resume across Transformer vs grouped HybridModel is still expected to hit config arg mismatches because Transformer uses logical num_layers=16, while grouped Hybrid uses physical symbols grouped into 16 logical layers.

[Wohox]

After removing cuda graph from these 2 methods, can we remove _apply_mamba_layer and _apply_attention_layer method and directly call layer.forward()?

[Connor-XY]

Both helpers are deleted and inlined at the single call site in pre_dispatch_computation (d2daf2925). Mamba uses item_layer(...) directly. Attention had to keep item_layer._forward_attention(...) instead of __call__, with this comment at the call site:

# Use _forward_attention rather than __call__: an attention half-layer has
# mlp=IdentityOp / mlp_bda=IdentityFuncOp by default, and TransformerLayer's
# __call__ would route through _forward_mlp + mlp_bda, double-applying the
# post-attention residual.

If we wanted to literally call __call__ for the attention case too, the attention layer spec would need mlp_bda set to something that returns mlp_output_with_bias[0] instead of the tuple. Happy to take that as a follow-up, but kept this round scoped to inlining.

[Wohox]

The current code works fine if pre_dispatch_computation starts with attn layer and ends with attn layer. We need to doublecheck:

• if pre_dispatch_computation starts with ssm/GDN, can we release the input tensor of pre_dispatch_computation after pre_dispatch_computation's forward pass .
• if pre_dispatch_computation ends with ssm/GDN, can we release the input tensor of dispatch after dispatch's forward pass.

Note: whether the input tensor can be released is determined by:

• whether the input tensor is needed by later submodule's forward, like the input tensor is a residual.
• whether the input tensor is added to saved_for_backward and can only be released after its backward pass finishes.

[Connor-XY]

Good catch. The current HybridStackNode._resolve_free_input delegates to the GPT default, which returns False for the pre_dispatch_computation slot — that's conservative but correct in every case I can reason about right now: the slot's input is the previous logical layer's output, and since hybrid groups always end the pre-dispatch phase by feeding hidden_states into MoE preprocess (which detaches the residual on node.layer_state.residual), nothing downstream of the slot actually retains a reference back to the slot's input.

The two cases you flag:

1. Pre-dispatch starts with SSM/GDN: the slot's input goes into Mamba/GDN forward, which produces a new hidden_states (Mamba does residual + mixer; GDN's _forward_attention returns the post-attn-bda hidden_states). Mamba/GDN backward needs the post-projection activations, not the original input — the input projection's saved-tensors handle that. So in principle the slot input could be freed.

2. Pre-dispatch ends with SSM/GDN (no terminal MoE): same reasoning for the final SSM/GDN layer's input. The dispatch slot is a no-op when there's no MoE terminal, so freeing the dispatch input is moot here.

I'd rather leave the policy at False for now and address this in a separate change once we've checked each backward path in detail (ideally with a memory measurement to confirm the saving). The seam to specialize is already there — _resolve_free_input is just a method override on HybridStackNode. Let me know if you'd prefer I make the change now and do the audit, or carry it as a follow-up.

[Wohox]

Sounds good, we can get back here as a follow-up.

[Connor-XY]

Filed mentally as a follow-up; will revisit when we audit memory peaks under SSM/GDN-only pre-dispatch.

[Wohox]

/claude review