Add the CSA/HCA prototype to HybridModel

examples/dsv4_hybrid/README.md

@@ -0,0 +1,94 @@
+# DeepSeek-V4 hybrid attention (CSA / HCA) in the hybrid model
+
+This directory holds an example slurm script for pretraining a Mamba-based hybrid
+model that mixes DeepSeek-V4's two new attention variants:
+
+* **CSA — Compressed Sparse Attention** (Section 2.3.1 of the DSv4 tech report).
+  Compresses every 4 tokens (with overlap = 2) into one KV entry and uses a
+  learned lightning indexer to pick the top-k relevant compressed entries. Each
+  query then attends to a small sliding window plus the selected compressed
+  positions.
+* **HCA — Heavily-Compressed Attention** (Section 2.3.2). Compresses every 128
+  tokens (no overlap) into one KV entry and applies dense attention over all
+  valid compressed positions, again concatenated with a sliding window.
+
+Both share the same module (`DSv4HybridSelfAttention` / `CompressedSparseAttention`);
+the per-layer compression ratio selects the behaviour.
+
+## Pattern symbols
+
+The hybrid layer pattern grows two new symbols in addition to the existing
+`M / G / * / D / - / E`:
+
+| Symbol | Meaning |
+| ------ | --------------------------------------------------------- |
+| `M`    | Mamba                                                     |
+| `G`    | Gated DeltaNet                                            |
+| `*`    | Standard self-attention (GQA)                             |
+| `D`    | DeepSeek Sparse Attention (DSA, MLA-style + indexer)      |
+| `C`    | DSv4 Compressed Sparse Attention (CSA, ratio = 4)         |
+| `H`    | DSv4 Heavily-Compressed Attention (HCA, ratio = 128)      |
+| `-`    | Dense MLP                                                 |
+| `E`    | MoE                                                       |
+
+Only one of `*` and the MLA-like family (`D`, `C`, `H`) may appear in the same
+model — they share the MLA-style q/kv projection setup.
+
+## Required CLI flags
+
+A model that uses `C` or `H` layers should set:
+
+* `--hybrid-layer-pattern` containing `C` and/or `H`.
+* The MLA-related flags: `--q-lora-rank`, `--qk-head-dim`, `--qk-pos-emb-head-dim`,
+  `--v-head-dim`, plus `--rope-type rope|yarn`.
+* DSA indexer flags (used only by `C`): `--dsa-indexer-n-heads`,
+  `--dsa-indexer-head-dim`, `--dsa-indexer-topk`,
+  optionally `--dsa-indexer-loss-coeff` for KL training of the indexer.
+
+CSA/HCA-specific flags:
+
+* `--csa-window-size N` — sliding-window length (default 128).
+* `--csa-compress-ratio-for-c N` — ratio used by every `C` layer (default 4).
+* `--csa-compress-ratio-for-h N` — ratio used by every `H` layer (default 128).
+* `--csa-compress-ratios "[...]"` — explicit per-layer ratios. Overrides the
+  pattern-derived defaults; length must equal `num_layers`.
+* `--csa-compress-rotary-base FLOAT` — RoPE base for compressed KV positions
+  (default 40000).
+* `--csa-dense-mode` — run all `C` layers in dense (no-indexer) mode. Useful as
+  a warmup phase before sparse training.
+* `--csa-no-attention-sink` — disable the per-head learnable sink logit.
+* `--o-groups N`, `--o-lora-rank N` — grouped output projection geometry.
+  ``num_attention_heads * v_head_dim`` must be divisible by `o_groups`.
+
+## Prototype scope
+
+This is a CP=1, TP=1 prototype that uses the unfused RoPE path. MTP, packed
+sequences, FP8/FP4, and fine-grained activation offloading are not supported
+yet. Inference is also disabled for `C`/`H` layers.
+
+## Running
+
+```
+bash examples/dsv4_hybrid/train_dsv4_hybrid.sh
+```
+
+The provided script trains a 2B-class hybrid model with the pattern
+
+```
+M-M-MCM-MHM-MDM-MCM-MHM-MDM-
+```
+
+(24 layers, mostly Mamba + MLP, with CSA / HCA / DSA layers sprinkled in).
+
+Adjust `IMAGE`, `BASE_DIR`, `BLEND_PATH`, `TOKENIZER_MODEL_PATH`, etc. to your
+environment. Tests must run inside the docker image — this is a slurm login
+node so GPUs are not directly available.
+
+## Tests
+
+* `tests/unit_tests/transformer/experimental_attention_variant/test_csa_hca_hybrid.py`
+  — unit tests covering the helpers, the `Compressor`, `CSAIndexer` and the
+  `CompressedSparseAttention` core attention (CSA, HCA and window-only paths,
+  forward + backward).
+* `tests/unit_tests/ssm/test_hybrid_layer_allocation.py`
+  — extended to exercise the new `C` and `H` symbols in the hybrid pattern.

examples/dsv4_hybrid/train_dsv4_hybrid.sh

@@ -0,0 +1,174 @@
+#!/bin/bash
+
+# DeepSeek-V4 hybrid model: Mamba + CSA / HCA / DSA + dense MLP pretraining example.
+#
+# Hybrid pattern symbols:
+#   M  - Mamba
+#   C  - DSv4 Compressed-Sparse Attention (CSA, ratio = 4 by default)
+#   H  - DSv4 Heavily-Compressed Attention (HCA, ratio = 128 by default)
+#   D  - DeepSeek Sparse Attention (DSA, MLA + indexer)
+#   -  - dense MLP
+#   E  - MoE
+#
+# Usage:
+#   ./dsv4_hybrid-1n.sh <checkpoint_path> <tensorboard_dir> <data_cache_dir> \
+#                               <tokenizer_model> <data_blend_path>
+
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export NVTE_FWD_LAYERNORM_SM_MARGIN=16
+export NVTE_BWD_LAYERNORM_SM_MARGIN=16
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+GPUS_PER_NODE=8
+MASTER_ADDR=localhost
+MASTER_PORT=6000
+NUM_NODES=1
+NODE_RANK=0
+WORLD_SIZE=$(($GPUS_PER_NODE*$NUM_NODES))
+
+CHECKPOINT_PATH=$1       # <Specify path>
+TENSORBOARD_LOGS_PATH=$2 # <Specify path>
+DATACACHE_PATH=$3        # <Specify path>
+TOKENIZER_MODEL=$4       # <Specify path to file>
+DATA_BLEND_PATH=$5       # <Specify path to data blend json>
+
+SEQ_LEN=8192
+TRAIN_SAMPLES=36621094
+LR_WARMUP_SAMPLES=1024000
+LR_DECAY_SAMPLES=36621094
+LR_WSD_DECAY_SAMPLES=5493165
+
+# 16-layer pattern: mamba bulk + interleaved CSA / HCA / DSA + dense MLP.
+HYBRID_PATTERN="M-MCM-MHMDM-MHM-"
+
+DISTRIBUTED_ARGS=(
+    --nproc_per_node $GPUS_PER_NODE
+    --nnodes $NUM_NODES
+    --master_addr $MASTER_ADDR
+    --master_port $MASTER_PORT
+)
+
+HYBRID_MODEL_ARGS=(
+    --hybrid-layer-pattern $HYBRID_PATTERN
+    --hidden-size 2048
+    --num-attention-heads 16
+    --num-query-groups 8
+    --ffn-hidden-size 8192
+    --kv-channels 128
+    --mamba-num-heads 64
+    --mamba-head-dim 64
+    --mamba-state-dim 128
+    --mamba-num-groups 8
+    --seq-length $SEQ_LEN
+    --max-position-embeddings $SEQ_LEN
+    --position-embedding-type none
+    --normalization RMSNorm
+    --untie-embeddings-and-output-weights
+    --disable-bias-linear
+    --squared-relu
+    --init-method-std 0.0198
+)
+
+# CSA / HCA / DSA arguments (apply to C, H, and D layers in the pattern).
+DSV4_HYBRID_ATTN_ARGS=(
+    --rope-type rope
+    --q-lora-rank 512
+    --qk-head-dim 128
+    --qk-pos-emb-head-dim 64
+    --v-head-dim 128
+    --csa-window-size 128
+    --csa-compress-ratio-for-c 4
+    --csa-compress-ratio-for-h 128
+    --o-groups 4
+    --o-lora-rank 128
+    --dsa-indexer-n-heads 64
+    --dsa-indexer-head-dim 128
+    --dsa-indexer-topk 64
+    --dsa-indexer-loss-coeff 0.0
+)
+
+TRAINING_ARGS=(
+    --micro-batch-size 1
+    --global-batch-size 8
+    --train-samples $TRAIN_SAMPLES
+    --lr 1.4e-3
+    --min-lr 1.4e-5
+    --lr-decay-style WSD
+    --lr-warmup-samples $LR_WARMUP_SAMPLES
+    --lr-decay-samples $LR_DECAY_SAMPLES
+    --lr-wsd-decay-style minus_sqrt
+    --lr-wsd-decay-samples $LR_WSD_DECAY_SAMPLES
+    --weight-decay 0.1
+    --adam-beta1 0.9
+    --adam-beta2 0.95
+    --clip-grad 1.0
+    --attention-dropout 0.0
+    --hidden-dropout 0.0
+    --bf16
+    --override-opt_param-scheduler
+)
+
+MODEL_PARALLEL_ARGS=(
+    --tensor-model-parallel-size 1
+    --pipeline-model-parallel-size 1
+    --use-distributed-optimizer
+    --overlap-grad-reduce
+    --overlap-param-gather
+)
+
+DATA_ARGS=(
+    --per-split-data-args-path $DATA_BLEND_PATH
+    --data-cache-path $DATACACHE_PATH
+    --tokenizer-type TikTokenizer
+    --tokenizer-model $TOKENIZER_MODEL
+    --tiktoken-pattern v2
+    --num-workers 1
+    --num-dataset-builder-threads 4
+    --no-create-attention-mask-in-dataloader
+    --no-mmap-bin-files
+)
+
+CHECKPOINT_ARGS=(
+    --save $CHECKPOINT_PATH
+    --load $CHECKPOINT_PATH
+    --ckpt-format torch_dist
+    --ckpt-fully-parallel-save
+    --ckpt-fully-parallel-load
+    --ckpt-assume-constant-structure
+    --no-load-rng
+    --save-interval 12500
+    --save-retain-interval 100000
+)
+
+EVAL_AND_LOGGING_ARGS=(
+    --log-interval 10
+    --log-throughput
+    --log-progress
+    --log-params-norm
+    --log-num-zeros-in-grad
+    --log-memory-interval 1000
+    --logging-level 20
+    --eval-interval 1000
+    --eval-iters 14
+    --tensorboard-dir $TENSORBOARD_LOGS_PATH
+)
+
+MISC_ARGS=(
+    --seed 1234
+    --rerun-mode disabled
+    --attention-backend flash
+    --disable-gloo-process-groups
+    --use-mcore-models
+    --spec megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec
+    --distributed-timeout-minutes 30
+)
+
+torchrun ${DISTRIBUTED_ARGS[@]} pretrain_hybrid.py \
+    ${HYBRID_MODEL_ARGS[@]} \
+    ${DSV4_HYBRID_ATTN_ARGS[@]} \
+    ${TRAINING_ARGS[@]} \
+    ${MODEL_PARALLEL_ARGS[@]} \
+    ${DATA_ARGS[@]} \
+    ${CHECKPOINT_ARGS[@]} \
+    ${EVAL_AND_LOGGING_ARGS[@]} \
+    ${MISC_ARGS[@]}

megatron/core/inference/contexts/dynamic_context.py

@@ -336,18 +336,41 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
             # For hybrid models, the layer map converts the global layer index to the
             # corresponding attention layer index or Mamba layer index depending on the
             # layer type.
-            attention_layer_map, dsa_layer_map, gdn_layer_map, mamba_layer_map = (
-                operator.itemgetter(
-                    Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.MAMBA
-                )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
+            (
+                attention_layer_map,
+                dsa_layer_map,
+                csa_layer_map,
+                hca_layer_map,
+                gdn_layer_map,
+                mamba_layer_map,
+            ) = operator.itemgetter(
+                Symbols.ATTENTION,
+                Symbols.DS_ATTENTION,
+                Symbols.CSA_ATTENTION,
+                Symbols.HCA_ATTENTION,
+                Symbols.GDN,
+                Symbols.MAMBA,
+            )(
+                get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)
             )
 
             if len(gdn_layer_map) > 0:
                 raise NotImplementedError("GDN layers are not supported for inference.")
 
-            self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map)
+            self.num_attention_layers = (
+                len(attention_layer_map)
+                + len(dsa_layer_map)
+                + len(csa_layer_map)
+                + len(hca_layer_map)
+            )
             self.num_mamba_layers = len(mamba_layer_map)
-            self.layer_map = attention_layer_map | dsa_layer_map | mamba_layer_map
+            self.layer_map = (
+                attention_layer_map
+                | dsa_layer_map
+                | csa_layer_map
+                | hca_layer_map
+                | mamba_layer_map
+            )
         else:
             # The layer map is the identity function for pure Transformer models.
             self.num_attention_layers = model_config.num_layers // pp_size

megatron/core/models/hybrid/hybrid_block.py

@@ -42,6 +42,8 @@ class HybridStackSubmodules:
     gdn_layer: Union[ModuleSpec, type] = IdentityOp
     attention_layer: Union[ModuleSpec, type] = IdentityOp
     dsa_layer: Union[ModuleSpec, type] = IdentityOp
+    csa_layer: Union[ModuleSpec, type] = IdentityOp
+    hca_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
     mtp_block_spec: Optional[ModuleSpec] = None
@@ -146,6 +148,24 @@ def __init__(
                         add_layer_offset=False,
                         pp_layer_offset=pp_layer_offset,
                     )
+                elif layer_type == LayerSymbols.CSA_ATTENTION:
+                    layer = build_module(
+                        submodules.csa_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        add_layer_offset=False,
+                        pp_layer_offset=pp_layer_offset,
+                    )
+                elif layer_type == LayerSymbols.HCA_ATTENTION:
+                    layer = build_module(
+                        submodules.hca_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        add_layer_offset=False,
+                        pp_layer_offset=pp_layer_offset,
+                    )
                 elif layer_type == LayerSymbols.MLP:
                     layer = build_module(
                         submodules.mlp_layer,

megatron/core/models/hybrid/hybrid_layer_allocation.py

@@ -18,11 +18,13 @@ class Symbols:
     GDN = 'G'
     ATTENTION = "*"
     DS_ATTENTION = "D"
+    CSA_ATTENTION = "C"
+    HCA_ATTENTION = "H"
     MLP = "-"
     MOE = 'E'
     PIPE = '|'
     MTP_SEPARATOR = "/"
-    VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLP, MOE}
+    VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, CSA_ATTENTION, HCA_ATTENTION, MLP, MOE}
 
     @classmethod
     def name_sorted_valid_layer_symbols(cls) -> list[str]:
@@ -292,9 +294,12 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False)
                 f"Valid symbols are: {valid_chars}"
             )
 
-    # Disallow Attention + MLA/DSA hybridity.
-    if Symbols.ATTENTION in pattern and Symbols.DS_ATTENTION in pattern:
-        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+    # Disallow Attention + MLA/DSA/CSA/HCA hybridity.
+    mla_like = {Symbols.DS_ATTENTION, Symbols.CSA_ATTENTION, Symbols.HCA_ATTENTION}
+    if Symbols.ATTENTION in pattern and any(s in pattern for s in mla_like):
+        raise ValueError(
+            "Not supported to have both standard Attention and MLA/DSA/CSA/HCA in one model"
+        )
 
 
 def validate_segment_layers(segment: str) -> List[str]:
@@ -320,9 +325,12 @@ def validate_segment_layers(segment: str) -> List[str]:
                 f"one of {Symbols.VALID_LAYERS}"
             )
 
-    # Disallow Attention + MLA/DSA hybridity.
-    if Symbols.ATTENTION in segment and Symbols.DS_ATTENTION in segment:
-        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+    # Disallow Attention + MLA/DSA/CSA/HCA hybridity.
+    mla_like = {Symbols.DS_ATTENTION, Symbols.CSA_ATTENTION, Symbols.HCA_ATTENTION}
+    if Symbols.ATTENTION in segment and any(s in segment for s in mla_like):
+        raise ValueError(
+            "Not supported to have both standard Attention and MLA/DSA/CSA/HCA in one model"
+        )
 
     return layer_type_list