|
| 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 | +* **mHC hyper-connections** (`ShensiHyperConnection`) — the residual stream is a pool of |
| 32 | + `hc_mult` streams. A sublayer collapses the pool into its input and writes its output back into |
| 33 | + a fixed plus a top-k routed subset of streams; the MLP side augments the write with causal depth |
| 34 | + convolutions and their Gram-Schmidt orthogonalizations. `ShensiHyperHead` collapses the pool at |
| 35 | + the end of the stack. |
| 36 | +* **Attention residuals** (`ShensiAttentionResidual`) — layers are grouped into blocks by the |
| 37 | + hash-MoE prefix and `attn_res_block_size`. A block-writing layer publishes its stream context and |
| 38 | + the later layers of the block read it. The history is per stream and per block, matching the |
| 39 | + reference's `[s, b, hc, nb, H]` buffer; `ShensiAttnResState` stores one tensor per block |
| 40 | + (`state.blocks`) and re-assembles that layout only for the pipeline payload and tooling, so a |
| 41 | + block boundary costs no buffer copy. The state also travels across pipeline stages on a |
| 42 | + dedicated communicator (the pipeline group's ranks, with its own tag space so the handoff's |
| 43 | + point-to-point ops cannot interleave with the schedule's), so a block may span stage boundaries. |
| 44 | + The state is per token and per stream (`hc_mult x num_blocks x hidden`), orders of magnitude |
| 45 | + larger than one activation, so for pipeline runs align the stage split with the |
| 46 | + attention-residual block boundaries — then nothing has to cross. The model logs the crossing |
| 47 | + blocks and that cost whenever a split does not line up. |
| 48 | +* **MoE layout** (`ShensiMoELayer`, `ShensiHashMLP`) — the leading layers are dense SwiGLU MLPs |
| 49 | + gated by a token-embedding lookup rather than routed experts, and the routed experts of the |
| 50 | + remaining layers run through a low-rank bottleneck whose norm is applied before the up |
| 51 | + projection. The ERC auxiliary loss needs a global view of the expert weights, which |
| 52 | + `erc_gather_expert_weights` builds across the expert-parallel group. |
| 53 | +* **Selective fp32 keeping** — the mHC mixers, the attention-residual readers and the strict list |
| 54 | + of norms / sinks / position biases stay in fp32 while the rest of the model trains in bf16, |
| 55 | + matching the `_keep_in_fp32_modules_strict` list of the HF checkpoint. |
| 56 | + |
| 57 | +## Field mapping |
| 58 | + |
| 59 | +The provider derives the Megatron-Core fields from a `ShensiConfig` through |
| 60 | +`shensi_derived_field_values`: `hc_mult → num_residual_streams`, `sliding_window → csa_window_size`, |
| 61 | +`compress_rope_theta → csa_compress_rotary_base`, `index_* → dsa_indexer_*`, |
| 62 | +`moe_intermediate_size → moe_ffn_hidden_size`, `routed_expert_hidden_size → moe_latent_size`, |
| 63 | +`o_groups / o_lora_rank → output_projection_groups / output_projection_lora_rank`. The HF config's |
| 64 | +`num_key_value_heads = 1` is expressed through `kv_lora_rank`, so `num_query_groups` follows the |
| 65 | +attention head count. Fields without a Megatron-Core counterpart are listed in |
| 66 | +`SHENSI_ONLY_FIELDS`; `describe_mapping()` prints the result for a built config. |
| 67 | + |
| 68 | +## Usage |
| 69 | + |
| 70 | +```python |
| 71 | +from megatron.bridge import AutoBridge |
| 72 | + |
| 73 | +bridge = AutoBridge.from_hf_pretrained("<hf checkpoint>") |
| 74 | +provider = bridge.to_megatron_provider(load_weights=False) |
| 75 | +model = provider.provide_distributed_model(wrap_with_ddp=False) |
| 76 | +bridge.load_hf_weights(model) |
| 77 | +``` |
| 78 | + |
| 79 | +## Inference |
| 80 | + |
| 81 | +Megatron-Core's hybrid sparse attention has no inference path (`DSv4HybridAttention` rejects |
| 82 | +inference contexts), so `examples/conversion/hf_to_megatron_generate_text.py` can convert and load |
| 83 | +a Shensi checkpoint but cannot generate with it. Serving is done through vLLM, which ships the |
| 84 | +attention and compressor kernels for this family: |
| 85 | + |
| 86 | +```sh |
| 87 | +vllm serve <exported hf directory> --served-model-name shensi --enforce-eager |
| 88 | +``` |
| 89 | + |
| 90 | +Exporting back to HuggingFace works (`bridge.save_hf_pretrained(model, path)`), which is what the |
| 91 | +functional round-trip test covers. |
| 92 | + |
| 93 | +## Tests |
| 94 | + |
| 95 | +```sh |
| 96 | +uv run python -m pytest -q tests/unit_tests/models/shensi --noconftest --import-mode=importlib |
| 97 | +uv run python -m pytest -q tests/functional_tests/test_groups/models/shensi |
| 98 | +``` |
| 99 | + |
| 100 | +The functional tests build a tiny Shensi checkpoint (`make_tiny_hf_model`) and cover the round trip |
| 101 | +(HF → Megatron → HF), the forward pass, right-padding behaviour and the provider/MTP configuration. |
| 102 | +The unit tests compare the attention-residual read bit-for-bit against the HF reference, pin the |
| 103 | +per-slot block layout to that reference (same draws, same rows, same gradients) and exercise the |
| 104 | +pipeline-stage handoff across two processes. |
0 commit comments