Repository navigation
feat(model): add Shensi model support - #6278
zongzhilou wants to merge 1 commit into
Conversation
cuichenx
added
feature
|
Hi @zongzhilou is there a public HF model for Shensi? i could not find it |
Hi! Shensi is our in-house architecture, so there's no public HF model yet. Merging a model PR into transformers requires the full weights to be open-sourced, but I'm still running experiments. Since Shensi was trained with mbridge, I've submitted the PR to mbridge first to add support for the architecture, and the two implementations are identical. Once the experiments are done, I'll release the weights and submit the PR to transformers. |
|
|
||
|
|
||
| @dataclass | ||
| class ShensiModelProvider(ShensiTransformerConfig, GPTModelProvider): |
There was a problem hiding this comment.
Not blocking: GPTModel is deprecated, and we recommend migrating to the HybridModel class.
Also, we recommend moving to the model config + model builder set up (see MuseGlimmerModelBuilder for an example)
There was a problem hiding this comment.
Done. The model rides HybridModel (ShensiModel, provider ShensiModelProvider(HybridModelProvider, ShensiTransformerConfig)), and shensi_builder.py adds the model-config/model-builder path
(ShensiModelConfig + ShensiModelBuilder(HybridModelBuilder)). Both delegate to
shensi_hybrid.build_shensi_model, so the provider and the builder construct the same model.
| def _hf_attention_residual(): | ||
| """The HF reference module (skips the test when the shensi HF implementation is absent).""" | ||
| pytest.importorskip("transformers.models.shensi.modeling_shensi") | ||
| from transformers.models.shensi.configuration_shensi import ShensiConfig |
There was a problem hiding this comment.
this will throw an error right now. Better guard it
There was a problem hiding this comment.
Guarded: pytest.importorskip for both shensi HF modules in
tests/unit_tests/models/shensi/test_shensi_attn_res_parity.py, plus skip-with-reason when
ShensiConfig(...) cannot be built — the test skips explicitly instead of failing or passing
vacuously.
| APPLY_GUARD_ATTR = "_shensi_orig_apply" | ||
| ENV_SWITCH = "SHENSI_HC_FP32_KEEP" | ||
| ENV_GROUPS = "SHENSI_FP32_KEEP_GROUPS" | ||
| ENV_NO_CAST = "SHENSI_FP32_KEEP_NO_CAST" |
There was a problem hiding this comment.
we usually don't use environment variables to control these settings; either hardcode one option or expose it properly if both options are needed.
There was a problem hiding this comment.
Removed, including the test-side ones (the old TE_FL_PREFER). Kernel selection belongs to the
launcher; grep -rn "environ" is empty across the shensi model package and its unit/functional
tests.
| return header, state | ||
|
|
||
|
|
||
| class ShensiAttnResPPHandoff: |
There was a problem hiding this comment.
this is fine, but just noting that upstream has NVIDIA/Megatron-LM#6840 as a draft, so this can be removed after 6840 lands
| if ranks is None: | ||
| return | ||
| if ranks not in _GROUP_CACHE: | ||
| _GROUP_CACHE[ranks] = torch.distributed.new_group(ranks=list(ranks)) |
There was a problem hiding this comment.
this is a little suspicious. can you check how MCore 6840 handles this without creating new process groups?
bbceb61 to
87ceabb
Compare
Shensi is a sparse-attention MoE family: every decoder layer picks one of three attention
paths from a per-layer plan (compressed sparse / heavily compressed / sliding window), the
residual stream is a pool of streams, layers group into attention-residual blocks that may
span pipeline stages, the leading layers are token-embedding-gated dense MLPs, and routed
experts use a low-rank bottleneck with an ERC auxiliary loss. The hybrid attention rides
Megatron-Core's dsv4_hybrid path.
- configuration_shensi: HF-side config plus the Megatron-Core transformer config (no mHC: the
residual path is the attention-residual read, so the V4 stream fields/knobs are not carried)
- modeling_shensi: attention residuals with pipeline-stage state
transfer (the block history is kept as one tensor per block, without a full-buffer copy at
every block boundary), MoE layout, selective fp32 keeping, decoder layer/stack and layer specs
- shensi_bridge / shensi_provider: config and parameter mapping, model provider
- shensi_builder: HybridModelConfig + HybridModelBuilder wiring so a training recipe can
build the model through Megatron-Core's model-config / model-builder API
- unit tests (attention-residual parity against the HF reference, per-slot block layout,
hybrid stack, provider, ERC gather, pipeline state transfer) and functional tests with an
L0 launcher
- model registry, README and docs entries
Follow-ups folded in while validating with a real 5-step smoke run:
- defaults follow the 2.78B / 1.24B-active geometry (routed_expert_hidden_size 320,
attn_res_block_size 2, o_groups 2, o_lora_rank 320, moe_intermediate_size 800, 1 MTP layer)
- attention-residual block writes are functional: the in-place variant bumped the storage
version of the buffer the read path holds in the autograd graph and failed in backward
- csa_dense_mode and the loss coefficients (indexer KL / aux / ERC) remain training-side
settings: they are no longer forced by hard defaults or by HF-side injection
- the provider path (AutoBridge conversion / provider smoke) builds again after the HybridModel
migration: it no longer reads the GPT-only `scatter_embedding_sequence_parallel`, and the MTP
depths map their inner decoder layer under `mtp.layers.{d}.mtp_model_layer.layers.0.` instead of
`mtp.layers.{d}.mtp_model_layer.`
- the MTP attention-residual state is bound on the inner decoder layers as well as on the wrapper
(the layers read it in forward, the wrapper is where the MTP layer and its test read it)
- `gradient_accumulation_fusion` re-declares the upstream auto-detect default on the provider, so a
host without apex fuses nothing instead of failing at model construction
- `finalize()` rejects Megatron-Core's own stream machinery (`enable_mhc_connections`,
`wide_residual`) the way Megatron-Core rejects an incompatible combination: shensi owns its
residual path inside the decoder layers, so the stock knobs would add a second one
- the functional tests no longer set `TE_FL_PREFER`; the model and test packages are env-free
- the attention-residual state crosses a pipeline stage boundary inside the stage payload itself
(`pack_attn_res_payload` / `unpack_attn_res_payload`, the Megatron-LM#6840 layout): the schedule
ships whatever the model returns, so the side-channel handoff, its private process group and its
header/serialization code are gone, and the state's gradients return through the same tensor; the
provider turns on `variable_seq_lengths` when the layout actually crosses, since the payload shape
then depends on the microbatch
The model carries no mHC: the attention-residual read is the residual path, so the V4
stream fields (`hc_mult` and friends), the mHC modules, their parameter mappings and the
per-stream state layout are gone (state, payload rows and fp32-keeping now describe one
residual stream).
- the hash-MoE deep embedding lives in the routed-expert bottleneck (a per-token scale
inside the low-rank space, RWKV-8 DeepEmbed's 4x form) and is applied between the
activation and the down projection, TP-sharded like the FC1 output
Signed-off-by: 313892219+zongzhilou <313892219+zongzhilou@users.noreply.github.com>
c7987b2 to
22fe063
Compare
This branch has not been deployed
2 participants
What does this PR do ?
Add Shensi model support to Megatron-Bridge: a sparse-attention Mixture-of-Experts family with its Hugging Face configuration, Megatron-Core provider, AutoBridge parameter mapping and tests.
Shensi is a sparse-attention Mixture-of-Experts family: every decoder layer picks one of three attention paths from a per-layer plan (compressed sparse attention, heavily compressed attention or a sliding window), the residual stream is a pool of streams that sublayers read from and write back into, layers are grouped into attention-residual blocks that may span pipeline stages, the leading layers are token-embedding-gated dense MLPs, and the routed experts of the remaining layers run through a low-rank bottleneck with an ERC auxiliary loss.
The hybrid attention itself is Megatron-Core's
dsv4_hybridpath; this bridge adds what sits around it and maps the checkpoint both ways.Changelog
tests/functional_tests/launch_scripts/h100/active/L0_Launch_models_shensi.shGitHub Actions CI
See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.
Before your PR is "Ready for review"
Pre checks:
If you haven't finished some of the above items you can still open "Draft" PR.
Additional Information
pytest tests/unit_tests/models/shensibash tests/functional_tests/launch_scripts/h100/active/L0_Launch_models_shensi.sh