|
| 1 | +# Shensi |
| 2 | + |
| 3 | +User-facing documentation lives under `docs/models/shensi/` (index page plus the variant page); |
| 4 | +this README covers the implementation. |
| 5 | + |
| 6 | +## Layout |
| 7 | + |
| 8 | +``` |
| 9 | +shensi/ |
| 10 | +├── configuration_shensi.py HF-side config: the per-layer plans, the RoPE split, validation |
| 11 | +├── modeling_shensi.py the components Shensi adds on top of Megatron-Core |
| 12 | +├── shensi_bridge.py HF <-> Megatron: config mapping and the parameter mappings |
| 13 | +├── shensi_hybrid.py HybridModel glue: pipeline layer pattern + hybrid stack spec |
| 14 | +├── shensi_builder.py model-config / model-builder path (HybridModelConfig + builder) |
| 15 | +├── shensi_provider.py provider: config fields -> ShensiModel, via build_shensi_model |
| 16 | +├── __init__.py |
| 17 | +└── README.md |
| 18 | +``` |
| 19 | + |
| 20 | +`configuration_shensi.py` holds both config surfaces: `ShensiConfig` mirrors the fields a |
| 21 | +checkpoint ships, `ShensiTransformerConfig` is the Megatron-Core transformer config that carries |
| 22 | +the Shensi-only fields next to the MLA ones. |
| 23 | + |
| 24 | +## What Shensi adds |
| 25 | + |
| 26 | +**Hybrid sparse attention** comes from Megatron-Core: the provider sets |
| 27 | +`experimental_attention_variant="dsv4_hybrid"` together with the per-layer `compress_ratios`, so |
| 28 | +CSA / HCA / sliding-window layers, their compressors and the Lightning-Indexer style router are the |
| 29 | +core implementations. This package adds what sits around them: |
| 30 | + |
| 31 | +* **Attention residuals** (`ShensiAttentionResidual`) — layers are grouped into blocks by the |
| 32 | + hash-MoE prefix and `attn_res_block_size`. A block-writing layer publishes its residual context |
| 33 | + and the later layers of the block read it. The history is per stream and per block, matching the |
| 34 | + reference's `[s, b, hc, nb, H]` buffer; `ShensiAttnResState` stores one tensor per block |
| 35 | + (`state.blocks`) and re-assembles that layout only for the pipeline payload and tooling, so a |
| 36 | + block boundary costs no buffer copy. When a block spans a pipeline stage boundary, the state |
| 37 | + rides the stage payload itself (`pack_attn_res_payload` / `unpack_attn_res_payload`): the |
| 38 | + schedule ships whatever the model returns, so no side communicator and no extra process group are |
| 39 | + involved, and the state's gradients travel back through the same tensor. The payload then carries |
| 40 | + the state per token (`(1 + num_blocks) x hidden` on top of the activations), so the schedule has to exchange shapes per microbatch |
| 41 | + (`variable_seq_lengths: true`) for a crossing split to run, which the layout check enforces. For |
| 42 | + pipeline runs, align the stage split with the attention-residual block boundaries — then nothing |
| 43 | + has to cross; the model logs the crossing blocks whenever a split does not line up. |
| 44 | +* **MoE layout** (`ShensiMoELayer`, `ShensiHashMLP`) — the leading layers are dense SwiGLU MLPs |
| 45 | + gated by a token-embedding lookup rather than routed experts, and the routed experts of the |
| 46 | + remaining layers run through a low-rank bottleneck whose norm is applied before the up |
| 47 | + projection. The ERC auxiliary loss needs a global view of the expert weights, which |
| 48 | + `erc_gather_expert_weights` builds across the expert-parallel group. |
| 49 | +* **Selective fp32 keeping** — the attention-residual readers and the strict list |
| 50 | + of norms / sinks / position biases stay in fp32 while the rest of the model trains in bf16, |
| 51 | + matching the `_keep_in_fp32_modules_strict` list of the HF checkpoint. |
| 52 | + |
| 53 | +## Field mapping |
| 54 | + |
| 55 | +The provider derives the Megatron-Core fields from a `ShensiConfig` through |
| 56 | +`shensi_derived_field_values`: `sliding_window → csa_window_size`, |
| 57 | +`compress_rope_theta → csa_compress_rotary_base`, `index_* → dsa_indexer_*`, |
| 58 | +`moe_intermediate_size → moe_ffn_hidden_size`, `routed_expert_hidden_size → moe_latent_size`, |
| 59 | +`o_groups / o_lora_rank → output_projection_groups / output_projection_lora_rank`. The HF config's |
| 60 | +`num_key_value_heads = 1` is expressed through `kv_lora_rank`, so `num_query_groups` follows the |
| 61 | +attention head count. Fields without a Megatron-Core counterpart are listed in |
| 62 | +`SHENSI_ONLY_FIELDS`; `describe_mapping()` prints the result for a built config. |
| 63 | + |
| 64 | +## Usage |
| 65 | + |
| 66 | +```python |
| 67 | +from megatron.bridge import AutoBridge |
| 68 | + |
| 69 | +bridge = AutoBridge.from_hf_pretrained("<hf checkpoint>") |
| 70 | +provider = bridge.to_megatron_provider(load_weights=False) |
| 71 | +model = provider.provide_distributed_model(wrap_with_ddp=False) |
| 72 | +bridge.load_hf_weights(model) |
| 73 | +``` |
| 74 | + |
| 75 | +## Inference |
| 76 | + |
| 77 | +Megatron-Core's hybrid sparse attention has no inference path (`DSv4HybridAttention` rejects |
| 78 | +inference contexts), so `examples/conversion/hf_to_megatron_generate_text.py` can convert and load |
| 79 | +a Shensi checkpoint but cannot generate with it. Serving is done through vLLM, which ships the |
| 80 | +attention and compressor kernels for this family: |
| 81 | + |
| 82 | +```sh |
| 83 | +vllm serve <exported hf directory> --served-model-name shensi --enforce-eager |
| 84 | +``` |
| 85 | + |
| 86 | +Exporting back to HuggingFace works (`bridge.save_hf_pretrained(model, path)`), which is what the |
| 87 | +functional round-trip test covers. |
| 88 | + |
| 89 | +## Tests |
| 90 | + |
| 91 | +```sh |
| 92 | +uv run python -m pytest -q tests/unit_tests/models/shensi --noconftest --import-mode=importlib |
| 93 | +uv run python -m pytest -q tests/functional_tests/test_groups/models/shensi |
| 94 | +``` |
| 95 | + |
| 96 | +The functional tests build a tiny Shensi checkpoint (`make_tiny_hf_model`) and cover the round trip |
| 97 | +(HF → Megatron → HF), the forward pass, right-padding behaviour and the provider/MTP configuration. |
| 98 | +The unit tests compare the attention-residual read bit-for-bit against the HF reference, pin the |
| 99 | +per-slot block layout to that reference (same draws, same rows, same gradients) and exercise the |
| 100 | +state carriage across pipeline stages (pack/unpack round trip, gradient path, per-stage conditions). |
0 commit comments