[Train] Add DeepSeek Engram (#1107)

**Add deepseek engram**, ref: [DeepSeek Engram
Paper](https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf),
[DeepSeek Engram Github](https://github.com/deepseek-ai/Engram)
- Support tensor parallel, pipeline parallel, sequence parallel,
distributed data parallel
- Support NgramHash Caching
- End to end training support
- CI/CD Tests
- CKPT Conversion: FlagScale to HuggingFace

TODO:
- Engram Embedding Offload 
- Engram Prefetch, attn/mlp computation and memory access overlapping
- FlagOS support, based on Megatron-LM-FL and TransformerEngine-FL

Author: lxd-cumt

.pre-commit-config.yaml

@@ -5,6 +5,7 @@ exclude: |
         flagscale/models/megatron/.*|
         tools/checkpoint/loader.*\.py|
         tools/checkpoint/saver.*\.py|
+        tools/checkpoint/qwen3_engram/modeling_hf/.*|
         examples/aquila/tokenizer*|
     )$
 repos:

examples/qwen3/conf/train/engram.yaml

@@ -0,0 +1,112 @@
+system:
+  no_shared_fs: ${experiment.runner.no_shared_fs}
+  num_workers: 2
+  tensor_model_parallel_size: 1
+  pipeline_model_parallel_size: 1
+  expert_model_parallel_size: 1
+  context_parallel_size: 1
+  sequence_parallel: true
+  use_distributed_optimizer: true
+  overlap_grad_reduce: true
+  overlap_param_gather: true
+  precision:
+    bf16: true
+    attention_softmax_in_fp32: true
+    accumulate_allreduce_grads_in_fp32: true
+  logging:
+    log_interval: 1
+    tensorboard_log_interval: 1
+    wandb_project: ${experiment.exp_name}
+    wandb_exp_name: ${experiment.exp_name}
+    log_timers_to_tensorboard: true
+    log_validation_ppl_to_tensorboard: true
+    log_throughput: true
+    log_params_norm: true
+    log_num_zeros_in_grad: true
+    log_memory_to_tensorboard: true
+  checkpoint:
+    save_interval: ${experiment.save_steps}
+    load: ${experiment.load}
+    ckpt_format: ${experiment.ckpt_format}
+
+model:
+  transformer_impl: transformer_engine
+  num_layers: 36
+  hidden_size: 2560
+  ffn_hidden_size: 9728
+  num_attention_heads: 32
+  kv_channels: 128
+  group_query_attention: true
+  num_query_groups: 8
+  seq_length: 4096
+  max_position_embeddings: 32768
+  norm_epsilon: 1e-6
+  use_rotary_position_embeddings: true
+  rotary_base: 10000
+  swiglu: true
+  normalization: RMSNorm
+  qk_layernorm: true
+  init_method_std: 0.02
+  attention_dropout: 0.0
+  hidden_dropout: 0.0
+  untie_embeddings_and_output_weights: true
+  position_embedding_type: rope
+  no_rope_fusion: true
+
+  # hf: attention_bias: true
+  disable_bias_linear: true
+  add_qkv_bias: false
+
+  # engram args =================
+  use_engram: true
+  engram_tokenizer_name_or_path: xxx
+  engram_vocab_size: [759680, 759680]
+  max_ngram_size: 3
+  n_embed_per_ngram: 512
+  n_head_per_ngram: 8
+  engram_layer_ids: [2]
+  engram_pad_id: 2
+  engram_seed: 0
+  engram_kernel_size: 4
+  engram_hc_mult: 1
+  # hf: engram_lr_multiplier: 5.0 miss
+
+
+  # training
+  seed: ${experiment.seed}
+  finetune: false
+  # to be update
+  micro_batch_size: 1
+  # hf: gbs: 1024
+  global_batch_size: 128
+  # hf: max_steps: 24000
+  train_iters: 102400
+  eval_iters: 0
+
+
+  optimizer:
+    clip_grad: 1.0
+    weight_decay: 0.1
+    adam_beta1: 0.9
+    adam_beta2: 0.95
+    lr_scheduler:
+      lr: 3.0e-4
+      min_lr: 3.0e-4
+      # lr_warmup_fraction: 0.01
+      lr_warmup_iters: 240
+      lr_decay_style: WSD
+      lr_wsd_decay_style: cosine
+      lr_wsd_decay_iters: 10
+
+data:
+  reset_position_ids: True
+  reset_attention_mask: True
+  data_path: xxx
+  split: 1
+  no_mmap_bin_files: true
+  tokenizer:
+    legacy_tokenizer: true
+    tokenizer_type: Qwen2TokenizerFS
+    tokenizer_path: xxx
+    vocab_size: 151936
+    make_vocab_size_divisible_by: 64

examples/qwen3/conf/train_engram.yaml

@@ -0,0 +1,33 @@
+defaults:
+  - _self_
+  - train: engram
+
+experiment:
+  exp_name: Qwen3-Engram
+  seed: 42
+  save_steps: 50
+  load: None
+  exp_dir: xxx
+  ckpt_format: torch
+  task:
+    type: train
+    backend: megatron
+    entrypoint: flagscale/train/megatron/train_engram.py
+  runner:
+    per_node_task: false
+    no_shared_fs: false
+    rdzv_backend: static
+    hostfile: xxx
+    ssh_port: xxx
+  cmds:
+    before_start: ulimit -n 1048576 && source /root/miniconda3/bin/activate flagscale-train
+  envs:
+    LOGLEVEL: "INFO"
+    CUDA_VISIBLE_DEVICES: "0,1,2,3,4,5,6,7"
+    CUDA_DEVICE_MAX_CONNECTIONS: 1
+
+action: run
+
+hydra:
+  run:
+    dir: ${experiment.exp_dir}/hydra

flagscale/models/megatron/engram/engram.py

@@ -0,0 +1,122 @@
+## built-in
+import copy
+import math
+
+## third-party
+import torch
+import torch.nn as nn
+
+from .engram_config import EngramConfig
+from .multi_head_embedding import MultiHeadEmbedding
+
+## engram
+from .ngram_hash import get_or_create_hash_mapping
+from .short_conv import ShortConv
+
+
+class Engram(nn.Module):
+    def __init__(self, engram_cfg: EngramConfig, layer_id):
+        super().__init__()
+        assert engram_cfg.engram_hc_mult == 1, (
+            "Engram do not support hyper-connection now, engram_hc_mult must be 1"
+        )
+        self.engram_cfg = engram_cfg
+        self.backbone_config = copy.deepcopy(engram_cfg)
+
+        self.layer_id = layer_id
+        global_hash_mapping = get_or_create_hash_mapping(
+            engram_vocab_size=engram_cfg.engram_vocab_size,
+            max_ngram_size=engram_cfg.max_ngram_size,
+            n_embed_per_ngram=engram_cfg.n_embed_per_ngram,
+            n_head_per_ngram=engram_cfg.n_head_per_ngram,
+            layer_ids=engram_cfg.engram_layer_ids,
+            tokenizer_name_or_path=engram_cfg.engram_tokenizer_name_or_path,
+            pad_id=engram_cfg.engram_pad_id,
+            seed=engram_cfg.engram_seed,
+        )
+        self.multi_head_embedding = MultiHeadEmbedding(
+            engram_cfg,
+            list_of_N=[
+                x for y in global_hash_mapping.vocab_size_across_layers[self.layer_id] for x in y
+            ],
+            D=engram_cfg.n_embed_per_ngram // engram_cfg.n_head_per_ngram,
+        )
+        self.short_conv = ShortConv(
+            hidden_size=self.backbone_config.hidden_size,
+            kernel_size=engram_cfg.engram_kernel_size,
+            dilation=engram_cfg.max_ngram_size,
+            hc_mult=self.backbone_config.engram_hc_mult,
+        )
+        engram_hidden_size = (engram_cfg.max_ngram_size - 1) * engram_cfg.n_embed_per_ngram
+        self.value_proj = nn.Linear(engram_hidden_size, self.backbone_config.hidden_size)
+        self.key_projs = nn.ModuleList(
+            [
+                nn.Linear(engram_hidden_size, self.backbone_config.hidden_size)
+                for _ in range(self.backbone_config.engram_hc_mult)
+            ]
+        )
+        self.norm1 = nn.ModuleList(
+            [
+                nn.RMSNorm(self.backbone_config.hidden_size)
+                for _ in range(self.backbone_config.engram_hc_mult)
+            ]
+        )
+        self.norm2 = nn.ModuleList(
+            [
+                nn.RMSNorm(self.backbone_config.hidden_size)
+                for _ in range(self.backbone_config.engram_hc_mult)
+            ]
+        )
+
+    def forward(self, hidden_states, hash_input_ids):
+        """
+        # hidden_states: [L, B, HC_MULT, D]
+        hidden_states: [L, B, D] # do not support hyper-connection now, hc_mult must be 1
+        input_ids: [B, L]
+
+        # return: [L, B, HC_MULT, D]
+        return: [L, B, D] # do not support hyper-connection now, hc_mult must be 1
+        """
+        assert hash_input_ids is not None, "Hash input ids can not be None for EngramModel"
+        # [B, L, N_GRAM * N_HEADS_PER_GRAM]
+        # fake hyper-connection
+        hidden_states = hidden_states.unsqueeze(2)
+
+        embeddings = self.multi_head_embedding(hash_input_ids).flatten(start_dim=-2)
+        # [L/tp_size, B, N_GRAM * N_HEADS_PER_GRAM, N_EMBED_PER_GRAM // N_HEADS_PER_GRAM]
+        # [L/tp_size, B, N_GRAM * N_EMBED_PER_NGRAM]
+
+        # Pre-compute scaling factor for efficiency
+        scale = 1.0 / math.sqrt(self.backbone_config.hidden_size)
+        gates = []
+        for hc_idx in range(self.backbone_config.engram_hc_mult):
+            key = self.key_projs[hc_idx](embeddings)
+            # [L/tp_size, B, HIDDEN_SIZE]
+            normed_key = self.norm1[hc_idx](key)
+
+            query = hidden_states[:, :, hc_idx, :]
+            # [L, B, HIDDEN_SIZE]
+            normed_query = self.norm2[hc_idx](query)
+
+            # Compute scaled dot product similarity
+            gate = torch.sum(normed_key * normed_query, dim=-1, keepdim=True) * scale
+            # Apply smooth absolute value transformation: sign(x) * sqrt(|x|)
+            # This is equivalent to: abs().clamp_min(1e-6).sqrt() * sign()
+            gate = torch.sign(gate) * torch.sqrt(torch.abs(gate).clamp_min(1e-6))
+            gate = torch.sigmoid(gate)
+            # [L, B, 1]
+
+            gates.append(gate)
+        gates = torch.stack(gates, dim=2)
+        # [L, B, HC_MULT, 1]
+
+        value = gates * self.value_proj(embeddings).unsqueeze(2)
+        # [L, B, HC_MULT, HIDDEN_SIZE]
+        output = value + self.short_conv(value)
+        # [L, B, HC_MULT, HIDDEN_SIZE]
+
+        # re-fake hyper-connection
+        assert output.shape[2] == 1, "Engram do not support hyper-connection now, hc_mult must be 1"
+        output = output.squeeze(2)
+
+        return output