diff --git a/docs/models/thinkingmachines/index.md b/docs/models/thinkingmachines/index.md index 4786cddb100..16c1ffaff09 100644 --- a/docs/models/thinkingmachines/index.md +++ b/docs/models/thinkingmachines/index.md @@ -2,7 +2,7 @@ title: Thinking Machines description: Miles recipes for Thinking Machines Lab models — Inkling (975 B), a multimodal MoE with short convolution, relative attention, and a shared-expert sink. --- -Miles ships a native Megatron recipe for **Inkling**, Thinking Machines Lab's 975 B / 41 B-active multimodal mixture-of-experts model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders. The same backend drives both full-parameter and LoRA RL. +Miles ships a native Megatron recipe for **Inkling**, Thinking Machines Lab's 975 B / 41 B-active multimodal mixture-of-experts model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders. ## Variants @@ -17,9 +17,9 @@ Inkling needs 16 nodes of 4× GB300 and the `radixark/miles:inkling` image: ```bash cd /root/miles -python scripts/run_inkling_975b.py train \ +python scripts/run_inkling.py train \ --model-name Inkling --train-mode full --task dapo_math \ --num-nodes 16 --num-gpus-per-node 4 ``` -See the [Inkling](/models/thinkingmachines/inkling) page for the architecture summary, HF → Megatron conversion, validated parallelism layouts, training attention backends, LoRA RL, and multimodal RL. +See the [Inkling](/models/thinkingmachines/inkling) page for the architecture summary, HF → Megatron conversion, validated parallelism layouts, training attention backends, and multimodal RL. diff --git a/docs/models/thinkingmachines/inkling-small.md b/docs/models/thinkingmachines/inkling-small.md index e0292a14d52..50a7f1ce292 100644 --- a/docs/models/thinkingmachines/inkling-small.md +++ b/docs/models/thinkingmachines/inkling-small.md @@ -5,7 +5,7 @@ description: Launch recipe for Inkling-Small (276 B), the compact sibling of Ink ## 1. Model Introduction -[Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small) is the compact member of Thinking Machines Lab's Inkling family: a 276 B-total / 12 B-active-parameter, 42-layer multimodal MoE (256 routed + 2 shared experts, top-6 sigmoid routing) that matches — and on some benchmarks exceeds — the flagship 975 B model (e.g. SWEBench Verified 80.2 vs 77.6) at a fraction of the deployment footprint. The architecture is the same as [Inkling](/models/thinkingmachines/inkling) — ShortConv, local/global relative attention, and the shared-expert-sink MoE — so everything on the Inkling page (architecture summary, attention backends, R3 routing replay, LoRA schema, multimodal RL) applies unchanged. This page only covers what differs: the model registry entry and the validated small-cluster launch profile. +[Inkling-Small](https://huggingface.co/thinkingmachines/Inkling-Small) is the compact member of Thinking Machines Lab's Inkling family: a 276 B-total / 12 B-active-parameter, 42-layer multimodal MoE (256 routed + 2 shared experts, top-6 sigmoid routing) that matches — and on some benchmarks exceeds — the flagship 975 B model (e.g. SWEBench Verified 80.2 vs 77.6) at a fraction of the deployment footprint. The architecture is the same as [Inkling](/models/thinkingmachines/inkling) — ShortConv, local/global relative attention, and the shared-expert-sink MoE — so everything on the Inkling page (architecture summary, attention backends, R3 routing replay, multimodal RL) applies unchanged. This page only covers what differs: the model registry entry and the validated small-cluster launch profile. ## 2. Supported Variants @@ -22,19 +22,13 @@ cd /root/miles # Full-parameter GRPO. 276 B fits with the CPU-offloaded optimizer - # no NVMe streaming needed (unlike the 975 B recipe). -python scripts/run_inkling_975b.py train \ - --model-name Inkling-Small --train-mode full --task dapo_math \ +python scripts/run_inkling.py train \ + --model-name Inkling-Small --task dapo_math \ --num-nodes 4 --num-gpus-per-node 8 \ + --lr 5e-5 --rollout-batch-size 64 --global-batch-size 128 \ --sglang-context-length 4096 --rollout-max-response-len 2048 \ --extra-args "--offload-train-target cpu --sglang-mem-fraction-static 0.65 \ --optimizer-cpu-offload --overlap-cpu-optimizer-d2h-h2d --use-precision-aware-optimizer" - -# LoRA GRPO (rank 32, all-linear), same cluster; adapter-only weight -# sync swaps in ~4 s per rollout. -python scripts/run_inkling_975b.py train \ - --model-name Inkling-Small --train-mode lora --task dapo_math \ - --num-nodes 4 --num-gpus-per-node 8 \ - --sglang-context-length 4096 --rollout-max-response-len 2048 ``` The model definition lives in `scripts/models/inkling-small.sh` (`MODEL_ARGS_NUM_LAYERS` overrides the layer count for sliced smoke/parity checkpoints). HF → `torch_dist` conversion uses the same tool as Inkling with this recipe file — a single 8-GPU node (TP8 EP8) converts it in one pass. @@ -45,4 +39,4 @@ The model definition lives in `scripts/models/inkling-small.sh` (`MODEL_ARGS_NUM |---|---|---|---|---|---|---|---| | H200 | 32 | 4 | on | 8 | 4 | 1 | `--decoder-last-pipeline-num-layers 7` (42 = 7×5 + 7) | -Batch shape is configurable from the launcher (`--rollout-batch-size`, `--global-batch-size`; defaults 32/64). The validated Small runs used 64/128 with `--lr 6e-6`, producing a steadily rising dapo-math reward curve; LoRA uses the launcher's `lr 5e-6` default. +Batch shape is configurable from the launcher (`--rollout-batch-size`, `--global-batch-size`; defaults 32/64). The validated Small runs used 64/128 with `--lr 5e-5`, producing a steadily rising dapo-math reward curve. diff --git a/docs/models/thinkingmachines/inkling.md b/docs/models/thinkingmachines/inkling.md index 3b4d49381bf..e79db4c716d 100644 --- a/docs/models/thinkingmachines/inkling.md +++ b/docs/models/thinkingmachines/inkling.md @@ -7,7 +7,7 @@ The complete Inkling RL implementation is open at the Miles pull request: [`radi ## 1. Model Introduction -[Inkling](https://huggingface.co/thinkingmachines/Inkling) is a mixture-of-experts transformer released by Thinking Machines Lab, with 975 B total parameters and 41 B active, a context window of up to 1 M tokens, and pretraining on 45 trillion tokens of text, images, audio and video. Its architecture introduces short convolution, attention with relative positional embedding, and a novel MoE design with a shared-expert sink. Miles implements Inkling as a native Megatron model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders, and the same backend drives both full-parameter and LoRA RL. +[Inkling](https://huggingface.co/thinkingmachines/Inkling) is a mixture-of-experts transformer released by Thinking Machines Lab, with 975 B total parameters and 41 B active, a context window of up to 1 M tokens, and pretraining on 45 trillion tokens of text, images, audio and video. Its architecture introduces short convolution, attention with relative positional embedding, and a novel MoE design with a shared-expert sink. Miles implements Inkling as a native Megatron model: local and global relative attention, the residual ShortConv, the shared-sink router and experts, and the image and audio encoders. **Key highlights**: @@ -32,13 +32,8 @@ docker pull radixark/miles:inkling # Full-parameter GRPO on 16 nodes x 4 GB300, inside the container cd /root/miles -python scripts/run_inkling_975b.py train \ - --model-name Inkling --train-mode full --task dapo_math \ - --num-nodes 16 --num-gpus-per-node 4 - -# LoRA GRPO (rank 32, all-linear), same cluster -python scripts/run_inkling_975b.py train \ - --model-name Inkling --train-mode lora --task dapo_math \ +python scripts/run_inkling.py train \ + --model-name Inkling --task dapo_math \ --num-nodes 16 --num-gpus-per-node 4 ``` @@ -72,12 +67,12 @@ Pass `--hf-checkpoint ` to the launcher when the weights are already on a ### 4.2 HF → Megatron `torch_dist` conversion -Inkling ships in BF16, so conversion is a single distributed `torch_dist` shard (no precision cast). The model definition comes from `scripts/models/inkling-975b.sh`: +Inkling ships in BF16, so conversion is a single distributed `torch_dist` shard (no precision cast). The model definition comes from `scripts/models/inkling.sh`: ```bash cd /root/miles -source scripts/models/inkling-975b.sh -PYTHONPATH=/root/Megatron-LM torchrun \ +source scripts/models/inkling.sh +CONVERT_KEEP_PP1=1 PYTHONPATH=/root/Megatron-LM torchrun \ --nproc-per-node 4 --nnodes 4 \ --master-addr ${MASTER_ADDR} --master-port 12345 \ --node-rank ${NODE_RANK} \ @@ -90,6 +85,8 @@ PYTHONPATH=/root/Megatron-LM torchrun \ --save /root/models/Inkling_torch_dist/ ``` +`CONVERT_KEEP_PP1=1` keeps the conversion at PP1: without it the converter auto-bumps PP toward the rank count, which is incompatible with `--tensor-model-parallel-size 4` at 16 ranks. + The saved `torch_dist` checkpoint is parallelism-agnostic: training can load it under any validated TP / PP / EP layout. Point the launcher at the result with `--torch-dist`. ### 4.3 Multi-node fan-out @@ -122,7 +119,7 @@ GRPO with truncated importance sampling. The launcher defaults: global batch siz ### 5.3 Training attention backends -`--inkling-attn-backend` selects the training-side attention implementation: +The `MILES_INKLING_ATTN_BACKEND` environment variable selects the training-side attention implementation: | Backend | Role | 8K packed, rel-extent 1024, GB300 | |---|---|---| @@ -142,23 +139,13 @@ SGLANG_ARGS=( --sglang-enable-multimodal --sglang-context-length 8192 - # full-parameter --sglang-mem-fraction-static 0.6 --sglang-max-running-requests 64 --sglang-max-total-tokens 327680 - - # LoRA (replaces the three flags above) - --sglang-ep-size 16 - --no-offload-rollout --no-offload-train - --sglang-mem-fraction-static 0.65 - --sglang-max-running-requests 32 - --sglang-lora-backend triton - --sglang-lora-use-virtual-experts - --sglang-max-loras-per-batch 1 ) ``` -Weight updates stream Megatron shards into SGLang's tensor layout in bounded buckets over CUDA IPC (colocated). In LoRA mode only the adapter is synchronized. +Weight updates stream Megatron shards into SGLang's tensor layout in bounded buckets over CUDA IPC (colocated). ### 5.5 Optimizer @@ -176,27 +163,14 @@ Weight updates stream Megatron shards into SGLang's tensor layout in bounded buc --micro-batch-size 1 ``` -Within a single GB300 rack the 975 B policy, gradients, and FP32 optimizer state exceed GPU memory, so Miles streams Megatron `DistributedOptimizer` state between a bounded GPU working set and node-local NVMe. The offload changes storage placement, not the update math. The paused training actor's weights are additionally disk-backed through `torch_memory_saver`. LoRA training skips both offloads and uses dynamic batching (`--use-dynamic-batch-size --max-tokens-per-gpu 4096`) instead of the fixed micro-batch. +Within a single GB300 rack the 975 B policy, gradients, and FP32 optimizer state exceed GPU memory, so Miles streams Megatron `DistributedOptimizer` state between a bounded GPU working set and node-local NVMe. The offload changes storage placement, not the update math. The paused training actor's weights are additionally disk-backed through `torch_memory_saver`. -1. **Fixed micro-batches for full-parameter runs.** Dynamic token packing exposes a PP-p2p × EP-all-to-all NCCL launch-order race on varlen shapes; the launcher pins `--micro-batch-size 1` for full-parameter training. Keep it pinned until the pad-to-fixed fix lands. -2. **One LoRA adapter per engine.** RL serving uses `--sglang-max-loras-per-batch 1` with the triton LoRA backend and virtual experts: the engine holds exactly the current policy's adapter (see [LoRA RL](#6-lora-rl)). +1. **Fixed micro-batches.** Dynamic token packing exposes a PP-p2p × EP-all-to-all NCCL launch-order race on varlen shapes; the launcher pins `--micro-batch-size 1`. Keep it pinned until the pad-to-fixed fix lands. -## 6. LoRA RL - -`--train-mode lora` switches the same backend and parallel stack to adapter-only training. The base model stays frozen in both runtimes, and GRPO updates only the low-rank factors: - -$$ -y = Wx + \frac{\alpha}{r}\,B(Ax). -$$ - -The adapter follows Inkling's released LoRA schema, covering the attention, dense MLP, MoE, and LM-head projections (defaults: rank 32, `alpha = 32`, all-linear). The routed experts use a shared-outer factorization: one factor is shared across experts while the expert-specific factors follow EP sharding. - -After each optimizer step, Miles exports a serving-ready adapter directly from the distributed training state and hands it to the colocated SGLang worker over CUDA IPC; the frozen base is never re-transferred. This cuts the weight update from 49.4 s to 2.5 s (20×) and brings training-step time to ~85 % of full-parameter training, with no optimizer offload needed. Released Inkling adapters in safetensors format support warm starts via `--lora-adapter-path`, and per-rank adapter checkpoints allow exact training resume. - -## 7. Multimodal RL +## 6. Multimodal RL -`--task geo3k` switches to the vision-language recipe: structured multimodal rollouts, Megatron-side vision and audio towers, and routing replay preserved across the media-expanded sequence. Each image or audio sentinel in the rendered prompt expands into its patch/frame positions before packing, and the R3 trace is indexed over the expanded sequence, so replay stays aligned even though SGLang made its routing decisions after expansion. Production recipes keep the vision and audio towers frozen and train the language model (or its adapter) around their embeddings; training the towers is available as an experimental option. The same path supports audio inputs and both full-parameter and LoRA training. +`--task geo3k` switches to the vision-language recipe: structured multimodal rollouts, Megatron-side vision and audio towers, and routing replay preserved across the media-expanded sequence. Each image or audio sentinel in the rendered prompt expands into its patch/frame positions before packing, and the R3 trace is indexed over the expanded sequence, so replay stays aligned even though SGLang made its routing decisions after expansion. Production recipes keep the vision and audio towers frozen and train the language model around their embeddings; training the towers is available as an experimental option. The same path supports audio inputs. diff --git a/miles/backends/megatron_utils/actor.py b/miles/backends/megatron_utils/actor.py index 059de956f09..667e96cb0c5 100644 --- a/miles/backends/megatron_utils/actor.py +++ b/miles/backends/megatron_utils/actor.py @@ -70,6 +70,12 @@ def _setup_disk_offload_reclaim(disk_dir: str) -> None: + """Wipe this rank's train disk-offload dir on startup and re-arm the atexit wipe. + + torch_memory_saver unlinks each backup file as its allocation is freed on a + graceful teardown, but a SIGKILL'd run leaves stale files behind. The dir is + per-rank (see actor_factory), so clearing it wholesale touches nobody else. + """ if not disk_dir: return shutil.rmtree(disk_dir, ignore_errors=True) diff --git a/miles/backends/megatron_utils/megatron_to_hf/__init__.py b/miles/backends/megatron_utils/megatron_to_hf/__init__.py index 0873e6f8e32..eef100777a7 100644 --- a/miles/backends/megatron_utils/megatron_to_hf/__init__.py +++ b/miles/backends/megatron_utils/megatron_to_hf/__init__.py @@ -2,6 +2,7 @@ from .deepseekv4 import convert_deepseekv4_to_hf from .glm4 import convert_glm4_to_hf from .glm4moe import convert_glm4moe_to_hf +from .inkling import convert_inkling_to_hf from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf from .llama import convert_llama_to_hf from .mimo import convert_mimo_to_hf @@ -52,6 +53,8 @@ def _convert_to_hf_core(args, model_name, name, param): converted_named_tensors = convert_qwen2_to_hf(args, name, param) elif "deepseekv4" in model_name: converted_named_tensors = convert_deepseekv4_to_hf(args, name, param) + elif "inkling" in model_name: + converted_named_tensors = convert_inkling_to_hf(args, name, param) elif "llama" in model_name: converted_named_tensors = convert_llama_to_hf(args, name, param) elif "mimo" in model_name: diff --git a/miles/backends/megatron_utils/megatron_to_hf/inkling.py b/miles/backends/megatron_utils/megatron_to_hf/inkling.py new file mode 100644 index 00000000000..f952ce267dc --- /dev/null +++ b/miles/backends/megatron_utils/megatron_to_hf/inkling.py @@ -0,0 +1,233 @@ +import json +import re +from functools import cache + +import torch + +from miles.backends.megatron_utils.misc_utils import strip_param_name_prefix + + +@cache +def _load_text_cfg(hf_checkpoint: str) -> dict: + with open(f"{hf_checkpoint}/config.json") as f: + cfg = json.load(f) + return cfg.get("text_config", cfg) + + +def _text_cfg(args) -> dict: + return _load_text_cfg(args.hf_checkpoint) + + +def _attn_geometry(args, layer_idx: int) -> tuple[int, int, int, int]: + """(nh, nkv, hd, d_rel) for the layer; SWA layers use the swa_* head counts.""" + t = _text_cfg(args) + is_local = layer_idx in set(t["local_layer_ids"]) + nh = t["swa_num_attention_heads"] if is_local else t["num_attention_heads"] + nkv = t["swa_num_key_value_heads"] if is_local else t["num_key_value_heads"] + hd = t["swa_head_dim"] if is_local else t["head_dim"] + return nh, nkv, hd, t["d_rel"] + + +class _PieceAccumulator: + """Collects multi-param groups (shared experts, gate halves) across convert calls.""" + + def __init__(self): + self._buckets: dict[tuple[str, int], dict[int, torch.Tensor]] = {} + + def put(self, kind: str, layer: int, sub_idx: int, tensor: torch.Tensor, expected: int): + """Stash one piece; return the ordered tensors once all `expected` arrived, else None.""" + bucket = self._buckets.setdefault((kind, layer), {}) + bucket[sub_idx] = tensor + if len(bucket) < expected: + return None + assert len(bucket) == expected, f"{kind} layer {layer}: got {len(bucket)} pieces, expected {expected}" + del self._buckets[(kind, layer)] + return [bucket[i] for i in sorted(bucket)] + + +_ACC = _PieceAccumulator() + + +def _qkv_tp_size(args): + try: + from miles.backends.training_utils.parallel import get_parallel_state + + return get_parallel_state().tp.size + except Exception: + return getattr(args, "tensor_model_parallel_size", 1) or 1 + + +def _qkv_blocks_from_gathered(args, layer_idx: int, param: torch.Tensor): + """Reconstruct full q|k|v|r from TP-all-gathered [q0|k0|v0|r0 | q1|k1|v1|r1 | ...] blocks.""" + nh, nkv, hd, d_rel = _attn_geometry(args, layer_idx) + full = [nh * hd, nkv * hd, nkv * hd, nh * d_rel] + full_rows = sum(full) + rows = param.shape[0] + + assert rows == full_rows, ( + f"qkvr layer {layer_idx}: gathered rows {rows} != full q|k|v|r rows {full_rows} " + f"(nh={nh} nkv={nkv} hd={hd} d_rel={d_rel}). all_gather_param should return the full " + f"(un-sharded) weight; got a partial/oversized tensor." + ) + + tp = _qkv_tp_size(args) + if tp == 1: + return tuple(param.split(full, dim=0)) + + assert full_rows % tp == 0, f"qkvr layer {layer_idx}: full rows {full_rows} not divisible by tp {tp}" + nh_l = nh // tp + nkv_l = max(1, nkv // tp) + per = [nh_l * hd, nkv_l * hd, nkv_l * hd, nh_l * d_rel] + assert sum(per) * tp == full_rows, ( + f"qkvr layer {layer_idx}: per-rank block {per} x tp {tp} = {sum(per)*tp} != full rows " + f"{full_rows}. Likely nkv({nkv}) < tp({tp}) (kv-head replication), which this de-shard " + f"does not handle -- the model keeps nkv>=tp at train time (global nkv=8, local nkv=16)." + ) + shards = param.split(sum(per), dim=0) + blk = [s.split(per, dim=0) for s in shards] + return tuple(torch.cat([blk[i][j] for i in range(tp)], dim=0) for j in range(4)) + + +def _reinterleave_w13(w13: torch.Tensor) -> torch.Tensor: + two_i = w13.shape[0] + assert two_i % 2 == 0, f"w13 first dim must be even, got {two_i}" + half = two_i // 2 + gate, up = w13[:half], w13[half:] + return torch.stack((gate, up), dim=1).reshape(w13.shape) + + +_GLOBAL_MAPPING = { + "embedding.word_embeddings.weight": "model.llm.embed.weight", + "embedding.embed_norm.weight": "model.llm.embed_norm.weight", + "decoder.final_layernorm.weight": "model.llm.norm.weight", + "output_layer.weight": "model.llm.unembed.weight", +} + + +def convert_inkling_to_hf(args, name, param): + mc = strip_param_name_prefix(name) + + if mc in _GLOBAL_MAPPING: + return [(_GLOBAL_MAPPING[mc], param)] + + m = re.match(r"decoder\.layers\.(\d+)\.(.+)", mc) + if m is None: + return [(mc, param)] + layer = int(m.group(1)) + rest = m.group(2) + hp = f"model.llm.layers.{layer}." + + if rest == "self_attention.linear_qkv.weight": + q, k, v, r = _qkv_blocks_from_gathered(args, layer, param) + return [ + (hp + "attn.wq_du.weight", q), + (hp + "attn.wk_dv.weight", k), + (hp + "attn.wv_dv.weight", v), + (hp + "attn.wr_du.weight", r), + ] + if rest == "self_attention.linear_qkv.layer_norm_weight": + return [(hp + "attn_norm.weight", param)] + if rest == "self_attention.linear_proj.weight": + return [(hp + "attn.wo_ud.weight", param)] + if rest == "self_attention.q_norm.weight": + return [(hp + "attn.q_norm.weight", param)] + if rest == "self_attention.k_norm.weight": + return [(hp + "attn.k_norm.weight", param)] + if rest == "self_attention.rel_proj": + return [(hp + "attn.rel_logits_proj.proj", param)] + if rest == "self_attention.k_sconv.weight": + return [(hp + "attn.k_sconv.weight", param)] + if rest == "self_attention.v_sconv.weight": + return [(hp + "attn.v_sconv.weight", param)] + if rest == "self_attention.attn_sconv.weight": + return [(hp + "attn_sconv.weight", param)] + + if rest == "pre_mlp_layernorm.weight": + return [(hp + "mlp_norm.weight", param)] + if rest == "mlp.mlp_sconv.weight": + return [(hp + "mlp_sconv.weight", param)] + + if rest == "mlp.linear_fc1.layer_norm_weight": + return [(hp + "mlp_norm.weight", param)] + if rest == "mlp.linear_fc1.weight": + return [(hp + "mlp.w13_dn.weight", _reinterleave_w13(param))] + if rest == "mlp.linear_fc2.weight": + return [(hp + "mlp.w2_md.weight", param)] + if rest == "mlp.global_scale": + return [(hp + "mlp.global_scale", param)] + + gate_sub = {"mlp.router.weight": 0, "mlp.router.shared_gate": 1}.get(rest) + if gate_sub is not None: + out = _ACC.put("gate", layer, gate_sub, param, 2) + if out is None: + return [] + return [(hp + "mlp.gate.weight", torch.cat(out, dim=0))] + if rest == "mlp.router.expert_bias": + return [(hp + "mlp.gate.bias", param)] + if rest == "mlp.router.global_scale": + return [(hp + "mlp.gate.global_scale", param)] + + em = re.match(r"mlp\.experts\.(linear_fc1|linear_fc2)\.weight(\d+)", rest) + if em is not None: + which, j = em.group(1), int(em.group(2)) + if which == "linear_fc1": + two_i = param.shape[0] + assert two_i % 2 == 0, ( + f"routed expert {j} layer {layer}: linear_fc1 rows {two_i} not even; " + f"expected contiguous [gate; up] (2*intermediate)" + ) + inter = _text_cfg(args).get("intermediate_size") + assert inter is None or two_i == 2 * inter, ( + f"routed expert {j} layer {layer}: linear_fc1 rows {two_i} != 2*intermediate_size " + f"({2 * inter}). The gathered expert is not the full contiguous [gate; up] this " + f"converter assumes (ETP must be 1; no TP split on the expert intermediate dim)." + ) + gate, up = param.chunk(2, dim=0) + return [ + (hp + f"mlp.experts.{j}.gate_proj.weight", gate), + (hp + f"mlp.experts.{j}.up_proj.weight", up), + ] + else: + return [(hp + f"mlp.experts.{j}.down_proj.weight", param)] + + sm = re.match(r"mlp\.shared_experts\.experts\.(\d+)\.(linear_fc1|linear_fc2)\.weight", rest) + if sm is not None: + e, which = int(sm.group(1)), sm.group(2) + ns = _text_cfg(args)["n_shared_experts"] + if which == "linear_fc1": + out = _ACC.put("shared_fc1", layer, e, _reinterleave_w13(param), ns) + if out is None: + return [] + return [(hp + "mlp.shared_experts.shared_w13_weight", torch.stack(out, dim=0))] + else: + out = _ACC.put("shared_fc2", layer, e, param, ns) + if out is None: + return [] + return [(hp + "mlp.shared_experts.shared_w2_weight", torch.stack(out, dim=0))] + + raise ValueError(f"convert_inkling_to_hf: unhandled param name {name!r} (rest={rest!r})") + + +def get_inkling_atomic_update_groups(args): + """Keep each accumulated group in one weight-sync bucket so the accumulator completes before flush.""" + from ..update_weight.common import AtomicUpdateGroup + + ns = _text_cfg(args)["n_shared_experts"] + + groups = [] + for which in ("linear_fc1", "linear_fc2"): + groups.append( + AtomicUpdateGroup( + key=f"inkling_shared_{which}", + suffixes=tuple(f".mlp.shared_experts.experts.{e}.{which}.weight" for e in range(ns)), + optional=True, + ) + ) + groups.append( + AtomicUpdateGroup( + key="inkling_gate", + suffixes=(".mlp.router.weight", ".mlp.router.shared_gate"), + optional=True, + ) + ) + return groups diff --git a/miles/backends/megatron_utils/model.py b/miles/backends/megatron_utils/model.py index 466fe473ca6..a58c5437d53 100644 --- a/miles/backends/megatron_utils/model.py +++ b/miles/backends/megatron_utils/model.py @@ -57,6 +57,11 @@ logger = logging.getLogger(__name__) +def _has_loadable_ckpt(load_dir: str | None) -> bool: + """Whether ``--load`` holds anything; ``load_checkpoint`` dispatches dist vs HF itself.""" + return bool(load_dir) and Path(load_dir).is_dir() and any(Path(load_dir).iterdir()) + + from .bridge_lora_helpers import _ensure_model_list, _setup_lora_model_via_bridge # noqa: F401 from .lora_utils import save_lora_checkpoint @@ -165,6 +170,12 @@ def setup_model_and_optimizer( config.timers = None if _is_muon_optimizer(config.optimizer): + if config.muon_split_qkv and "inkling" in (getattr(args, "custom_model_provider_path", None) or ""): + if is_first_replica_megatron_main_rank(): + logger.info( + "Inkling fused qkvr detected: forcing muon_split_qkv=False " "(whole-matrix orthogonalization)." + ) + config.muon_split_qkv = False optimizer = get_megatron_muon_optimizer( config=config, model_chunks=model, @@ -957,8 +968,7 @@ def initialize_model_and_optimizer( model[0].role = role clear_memory() - multi_lora = is_multi_lora_enabled(args) - if multi_lora: + if is_multi_lora_enabled(args): # Hide adapter params so the bridge's conversion-task walk doesn't see them # while loading the base checkpoint. from megatron.bridge.peft.multi_lora_layers import hide_adapters @@ -967,14 +977,22 @@ def initialize_model_and_optimizer( else: load_ctx = nullcontext() - with load_ctx: - iteration, _ = load_checkpoint( - model, - optimizer, - opt_param_scheduler, - checkpointing_context=checkpointing_context, - skip_load_to_model_and_opt=False, - ) + load_dir = getattr(args, "load", None) + # --load may be unset: setup_model_and_optimizer already asserted pretrained_checkpoint covers it. + if load_dir is None or _has_loadable_ckpt(load_dir): + with load_ctx: + iteration, _ = load_checkpoint( + model, + optimizer, + opt_param_scheduler, + checkpointing_context=checkpointing_context, + skip_load_to_model_and_opt=False, + ) + else: + if is_first_replica_megatron_main_rank(): + logger.warning("--load %r is empty; starting from model_provider-initialized weights", load_dir) + iteration = 0 + check_peak_gpu_memory_after_load(args) clear_memory() diff --git a/miles/backends/megatron_utils/update_weight/common.py b/miles/backends/megatron_utils/update_weight/common.py index e003e4ae0ba..40aca036858 100644 --- a/miles/backends/megatron_utils/update_weight/common.py +++ b/miles/backends/megatron_utils/update_weight/common.py @@ -22,10 +22,11 @@ class AtomicUpdateGroup: key: str suffixes: tuple[str, ...] + optional: bool = False def get_atomic_update_groups(args, model_name) -> list[AtomicUpdateGroup]: - model_groups = _get_model_atomic_update_groups(model_name) + model_groups = _get_model_atomic_update_groups(args, model_name) if model_groups: return model_groups return _get_q_lora_atomic_update_groups(args) @@ -46,8 +47,12 @@ def _get_q_lora_atomic_update_groups(args) -> list[AtomicUpdateGroup]: ] -def _get_model_atomic_update_groups(model_name) -> list[AtomicUpdateGroup]: +def _get_model_atomic_update_groups(args, model_name) -> list[AtomicUpdateGroup]: model_name = model_name.lower() + if "inkling" in model_name: + from ..megatron_to_hf.inkling import get_inkling_atomic_update_groups + + return get_inkling_atomic_update_groups(args) if "deepseekv4" in model_name: from ..megatron_to_hf.deepseekv4 import get_deepseek_v4_atomic_update_groups @@ -103,7 +108,7 @@ def get_named_update_units(param_names: Sequence[str], atomic_update_groups) -> for group in atomic_update_groups: assert ( - group.key in matched_group_keys + group.optional or group.key in matched_group_keys ), f"Atomic update group {group.key} references no params matching suffixes {group.suffixes}" for (prefix, key), names in pending_groups.items(): @@ -138,13 +143,23 @@ def _gather_with_stride( return torch.cat(interleaved, dim=partition_dim) +def is_routed_expert_param(name: str) -> bool: + """Whether a Megatron param name belongs to the routed (expert-parallel) experts. + + Routed experts live under ".experts.", but shared experts may nest an inner + ModuleList (e.g. Inkling's "mlp.shared_experts.experts.N.") whose params are + regular-TP sharded and EP-replicated, so they must not match. + """ + return ".experts." in name and ".shared_experts." not in name + + def _is_unmarked_grouped_expert_weight(name: str, param: torch.nn.Parameter) -> bool: """TEGroupedLinear never marks its per-expert weight0..weightN, so Megatron fills in the defaults (tensor_model_parallel=False, partition_dim=-1) and the tensor claims to be unsharded. It is expert-TP sharded whenever etp > 1, so the gather must still run. """ return ( - ".experts." in name + is_routed_expert_param(name) and ("linear_fc1.weight" in name or "linear_fc2.weight" in name) and not param.tensor_model_parallel and get_parallel_state().etp.size > 1 @@ -174,7 +189,8 @@ def _check_and_fix_partition(args: Namespace, name: str, partition_stride: int, def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> torch.Tensor: """ All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data. - Uses expert-TP for ".experts.", else regular-TP. Handles strided partitioning via partition_stride. + Uses expert-TP for routed ".experts." (but NOT ".shared_experts." — those are split on the + REGULAR TP group like attention), else regular-TP. Handles strided partitioning via partition_stride. """ if "expert_bias" in name: return param @@ -185,13 +201,16 @@ def all_gather_param(args: Namespace, name: str, param: torch.nn.Parameter) -> t if not param.tensor_model_parallel and not _is_unmarked_grouped_expert_weight(name, param): return param.data - if ".experts." in name: + if is_routed_expert_param(name): tp_size = get_parallel_state().etp.size tp_group = get_parallel_state().etp.group else: tp_size = get_parallel_state().tp.size tp_group = get_parallel_state().tp.group + if tp_size <= 1: + return param.data + param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] dist.all_gather(param_partitions, param.data, group=tp_group) partition_dim = param.partition_dim @@ -227,13 +246,18 @@ def all_gather_params_async( handles.append(None) else: # Start async all_gather - if ".experts." in info.name: + if is_routed_expert_param(info.name): tp_size = get_parallel_state().etp.size tp_group = get_parallel_state().etp.group else: tp_size = get_parallel_state().tp.size tp_group = get_parallel_state().tp.group + if tp_size <= 1: + gather_tasks.append((info, param.data, None, None, None, None)) + handles.append(None) + continue + param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)] handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True) gather_tasks.append((info, None, handle, param_partitions, param.partition_dim, param.partition_stride)) @@ -402,7 +426,7 @@ def collect_named_tensors_for_weight_transfer( convert_to_global_name, translate_gpu_to_cpu, ): - if is_expert is None or is_expert == (".experts." in name): + if is_expert is None or is_expert == is_routed_expert_param(name): yield name, tensor diff --git a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py index 752ea491b0a..292c28e5e63 100644 --- a/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py +++ b/miles/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py @@ -17,6 +17,7 @@ all_gather_params_async, get_atomic_update_groups, get_named_update_units, + is_routed_expert_param, named_params_and_buffers, ) from .hf_weight_iterator_base import HfWeightIteratorBase @@ -91,7 +92,7 @@ def _get_megatron_full_params( if ep_size > 1: handles = [] for info, param in zip(megatron_local_param_infos, params, strict=False): - if ".experts." in info.name: + if is_routed_expert_param(info.name): src_rank = ( info.src_rank if info.src_rank in dist.get_process_group_ranks(get_parallel_state().ep.group) @@ -133,7 +134,7 @@ def _get_megatron_local_param_info_buckets( def _get_param_full_size(info: ParamInfo) -> int: - if ".experts." in info.name: + if is_routed_expert_param(info.name): tp_size = get_parallel_state().etp.size else: tp_size = get_parallel_state().tp.size diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py index b722a6164ba..fc1d01eed17 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_distributed/mixin.py @@ -21,6 +21,7 @@ end_weight_update, get_atomic_update_groups, get_named_value_update_units, + is_routed_expert_param, weight_update_selector, ) from ..hf_weight_iterator_base import HfWeightIteratorBase @@ -31,7 +32,7 @@ def _is_expert_update_unit(update_unit: list[tuple[str, torch.Tensor]]) -> bool: assert update_unit, "Update unit must contain at least one param" name, _tensor = update_unit[0] - return ".experts." in name + return is_routed_expert_param(name) class DistBucketedWeightUpdateMixin: @@ -177,7 +178,7 @@ def _get_weight_transfer_update_units(self, is_expert: bool) -> list[list[tuple[ atomic_update_groups = get_atomic_update_groups(self.args, self.model_name) update_units = get_named_value_update_units(named_tensors, atomic_update_groups) for unit in update_units: - assert len({".experts." in name for name, _tensor in unit}) == 1, [name for name, _tensor in unit] + assert len({is_routed_expert_param(name) for name, _tensor in unit}) == 1, [name for name, _tensor in unit] return [unit for unit in update_units if _is_expert_update_unit(unit) == is_expert] def _update_expert_bucket_weights( diff --git a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py index 755e75c0e95..17ba25bcbbf 100644 --- a/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py +++ b/miles/backends/megatron_utils/update_weight/update_weight_from_tensor.py @@ -1,5 +1,6 @@ import hashlib import logging +import os from argparse import Namespace from collections.abc import Callable, Mapping, Sequence from typing import Any @@ -71,9 +72,10 @@ def __init__( self._lora_loaded = False self._lora_base_synced = False - # Create IPC gather groups within megatron. + self._mm_tower_cache: list[tuple[str, torch.Tensor]] | None = None + for start_rank in range(0, dist.get_world_size(), self.args.rollout_num_gpus_per_engine): - end_rank = start_rank + self.args.rollout_num_gpus_per_engine + end_rank = min(start_rank + self.args.rollout_num_gpus_per_engine, dist.get_world_size()) group_ranks = list(range(start_rank, end_rank)) new_group = dist.new_group(ranks=group_ranks, backend="gloo") if dist.get_rank() in group_ranks: @@ -229,6 +231,16 @@ def update_weights(self) -> None: _check_weight_sync_results(results, is_lora=False) del long_lived_tensors + mm_tower_tensors = self._mm_tower_named_tensors() + if mm_tower_tensors is not None: + mm_tower_tensors = [ + (name, tensor.to(torch.cuda.current_device())) for name, tensor in mm_tower_tensors + ] + refs, long_lived_tensors = self._send_base_params(mm_tower_tensors) + results = ray.get(refs) + _check_weight_sync_results(results, is_lora=False) + del long_lived_tensors, mm_tower_tensors + if self.is_lora: # SGLang's load_lora_adapter_from_tensors expects the full adapter in # one call; drain the bridge's chunker so --update-weight-buffer-size @@ -263,6 +275,50 @@ def update_weights(self) -> None: ray.get([engine.continue_generation.remote() for engine in self.rollout_engines]) dist.barrier(group=get_gloo_group()) + def _mm_tower_named_tensors(self) -> list[tuple[str, torch.Tensor]] | None: + """Frozen vision/audio tower tensors to append to every base sync (see + __init__ comment). Returns None when the run has no MM towers. EVERY + gather-group rank contributes the full tower set (read once from its local + HF checkpoint, the same bytes the engine loaded at boot): the colocated + send requires homogeneous per-rank bucket counts (num_dtypes is taken from + rank 0 and indexed into every rank's list), so a src-only contribution + breaks assembly. The duplicates are ~15MB/rank and load idempotently.""" + provider = getattr(self.args, "custom_model_provider_path", None) or "" + if "inkling_mm_model_provider" not in provider: + return None + if self._mm_tower_cache is None: + if self._ipc_gather_group is not None: + import json + + from safetensors import safe_open + + ckpt_dir = self.args.hf_checkpoint + with open(os.path.join(ckpt_dir, "model.safetensors.index.json"), encoding="utf-8") as f: + weight_map = json.load(f)["weight_map"] + tower_keys = sorted( + k + for k in weight_map + if ".visual." in f".{k}" or ".audio." in f".{k}" or k.startswith(("visual.", "audio.")) + ) + by_shard: dict[str, list[str]] = {} + for k in tower_keys: + by_shard.setdefault(weight_map[k], []).append(k) + cache = [] + for shard, keys in by_shard.items(): + with safe_open(os.path.join(ckpt_dir, shard), framework="pt", device="cpu") as f: + for k in keys: + cache.append((k, f.get_tensor(k))) + logger.info( + "mm tower sync: caching %d tower tensors from %s: %s", + len(cache), + ckpt_dir, + [k for k, _ in cache], + ) + self._mm_tower_cache = cache + else: + self._mm_tower_cache = [] + return self._mm_tower_cache + def _send_base_params(self, hf_named_tensors) -> tuple[list[ObjectRef], Any]: refs, long_lived_tensors = _send_to_colocated_engine( hf_named_tensors=hf_named_tensors, diff --git a/miles/backends/sglang_utils/sglang_engine.py b/miles/backends/sglang_utils/sglang_engine.py index d61b01b1fd7..c19dce207e6 100644 --- a/miles/backends/sglang_utils/sglang_engine.py +++ b/miles/backends/sglang_utils/sglang_engine.py @@ -191,7 +191,6 @@ def _format_v6_uri(addr): host = _format_v6_uri(host) ip_part, port_part = dist_init_addr.rsplit(":", 1) dist_init_addr = f"{_format_v6_uri(ip_part)}:{port_part}" - server_args_dict, external_engine_need_check_fields = _compute_server_args( self.args, self.rank, @@ -727,6 +726,9 @@ def _compute_server_args( "enable_metrics": True, } + if os.environ.get("MILES_SGLANG_DUMMY_LOAD") == "1": + kwargs["load_format"] = "dummy" + if sglang_overrides: kwargs.update(sglang_overrides) diff --git a/miles/backends/training_utils/data.py b/miles/backends/training_utils/data.py index 9f96e2884fa..3dce74be5d6 100644 --- a/miles/backends/training_utils/data.py +++ b/miles/backends/training_utils/data.py @@ -324,11 +324,16 @@ def _compute_transform_like_token_ids(ids_list: list): # Process multimodal training tensors if present multimodal_train_inputs = batch.get("multimodal_train_inputs", None) if multimodal_train_inputs is not None: + sample_offsets = [0] + for t in batch["unconcat_tokens"]: + sample_offsets.append(sample_offsets[-1] + t.size(0)) multimodal_data = {} # key -> concatenated tensor multimodal_num_items = {} # key -> list of item counts per sequence - for mm_input_dict in multimodal_train_inputs: + for i, mm_input_dict in enumerate(multimodal_train_inputs): if mm_input_dict is not None: for key, mm_tensor in mm_input_dict.items(): + if key.endswith("_positions"): + mm_tensor = mm_tensor + sample_offsets[i] if key not in multimodal_data: multimodal_data[key] = mm_tensor multimodal_num_items[key] = [mm_tensor.size(0)] diff --git a/miles/backends/training_utils/mm_data.py b/miles/backends/training_utils/mm_data.py index c19ca7ba0fc..c0f56d7b81f 100644 --- a/miles/backends/training_utils/mm_data.py +++ b/miles/backends/training_utils/mm_data.py @@ -1,10 +1,11 @@ """Multimodal-specific preprocessing for rollout data. -Currently Kimi-VL / Kimi-K2.5-only: the rollout side emits one media -placeholder token per image, and training expands it to grid-derived token -counts so the LM sees a position per vision patch. Kept separate from -data.py (generic batching / CP slicing) so additional VL models can land -without inflating that module. +The rollout side emits one media placeholder/sentinel per media item; training +expands it to the per-item token count so the LM sees a position per vision +patch / audio frame. Two families live here: Kimi-VL / Kimi-K2.5 (grid-derived +expansion of an in-vocab placeholder) and Inkling (out-of-vocab sentinels expanded +to in-vocab placeholder runs with explicit positions). Kept separate from +data.py (generic batching / CP slicing). """ import logging @@ -114,12 +115,103 @@ def _batch_has_media_placeholders( return any((token_tensor == media_token_id).any().item() for token_tensor in tokens) +INKLING_IMAGE_SENTINEL_ID = -101 +INKLING_AUDIO_SENTINEL_ID = -102 +INKLING_MM_PLACEHOLDER_TOKEN_ID = 200023 +INKLING_MM_AUDIO_PLACEHOLDER_TOKEN_ID = 200025 + + +def _expand_inkling_sample(token_tensor, loss_mask, mm, sample_idx: int): + """Replace media sentinels with placeholder runs, recording their sample-local positions into mm; idempotent.""" + spec = ( + (INKLING_IMAGE_SENTINEL_ID, INKLING_MM_PLACEHOLDER_TOKEN_ID, "mm_vision_num_patches", "mm_vision_positions"), + ( + INKLING_AUDIO_SENTINEL_ID, + INKLING_MM_AUDIO_PLACEHOLDER_TOKEN_ID, + "mm_audio_num_tokens", + "mm_audio_positions", + ), + ) + splices = [] + for sentinel, placeholder, counts_key, positions_key in spec: + counts = mm.get(counts_key) if mm else None + positions = (token_tensor == sentinel).nonzero(as_tuple=True)[0] + if positions.numel() == 0: + continue + assert counts is not None and positions.numel() == len(counts), ( + f"sample {sample_idx}: {positions.numel()} sentinel(s) {sentinel} but " + f"{counts_key}={'missing' if counts is None else len(counts)}" + ) + splices.append((positions.tolist(), [int(c) for c in counts], placeholder, positions_key)) + + if not splices: + return token_tensor + + prompt_len = len(token_tensor) - len(loss_mask) + flat = sorted( + (pos, n, placeholder, positions_key) + for positions, counts, placeholder, positions_key in splices + for pos, n in zip(positions, counts, strict=True) + ) + assert all( + p < prompt_len for p, _, _, _ in flat + ), "Inkling media sentinels must be in the prompt; found one in the response" + + pieces, prev, shift = [], 0, 0 + out_positions: dict[str, list[int]] = {} + for pos, n, placeholder, positions_key in flat: + pieces.append(token_tensor[prev:pos]) + pieces.append(torch.full((n,), placeholder, dtype=token_tensor.dtype, device=token_tensor.device)) + start = pos + shift + out_positions.setdefault(positions_key, []).extend(range(start, start + n)) + shift += n - 1 + prev = pos + 1 + pieces.append(token_tensor[prev:]) + for positions_key, plist in out_positions.items(): + mm[positions_key] = torch.tensor(plist, dtype=torch.long, device=token_tensor.device) + return torch.cat(pieces) + + +def _expand_inkling_rollout_data_in_place(rollout_data: RolloutBatch) -> None: + mm_list = rollout_data["multimodal_train_inputs"] + tokens = rollout_data["tokens"] + loss_masks = rollout_data["loss_masks"] + old_total_lengths = list(rollout_data["total_lengths"]) + + new_tokens_list, new_total_lengths = [], [] + for i, (token_tensor, loss_mask) in enumerate(zip(tokens, loss_masks, strict=False)): + mm = mm_list[i] if i < len(mm_list) else None + expanded = _expand_inkling_sample(token_tensor, loss_mask, mm, i) + new_tokens_list.append(expanded) + new_total_lengths.append(expanded.size(0)) + + if new_total_lengths != old_total_lengths: + try: + cp_size = get_parallel_state().cp.size + except Exception: + cp_size = 1 + assert cp_size == 1, "Inkling multimodal expansion does not support CP>1 yet" + rollout_data["tokens"] = new_tokens_list + rollout_data["total_lengths"] = new_total_lengths + logger.info( + "Expanded Inkling image sentinels: total_lengths %s -> %s", + old_total_lengths, + new_total_lengths, + ) + + def expand_multimodal_rollout_data_in_place( rollout_data: RolloutBatch, media_token_id: int = KIMI_VL_MEDIA_TOKEN_ID, qkv_format: str = "thd", ) -> None: multimodal_train_inputs = rollout_data.get("multimodal_train_inputs", None) + if multimodal_train_inputs is not None and any( + mm is not None and ("mm_vision_num_patches" in mm or "mm_audio_num_tokens" in mm) + for mm in multimodal_train_inputs + ): + _expand_inkling_rollout_data_in_place(rollout_data) + return mm_inputs_list = _collect_multimodal_grid_inputs(multimodal_train_inputs) if not mm_inputs_list or not any(mm is not None for mm in mm_inputs_list): return diff --git a/miles/rollout/generate_utils/generate_endpoint_utils.py b/miles/rollout/generate_utils/generate_endpoint_utils.py index ac90f9fc2d0..064254848db 100644 --- a/miles/rollout/generate_utils/generate_endpoint_utils.py +++ b/miles/rollout/generate_utils/generate_endpoint_utils.py @@ -125,7 +125,7 @@ def get_routed_experts_from_response(args, output, sample): info = output["meta_info"].get("routed_experts") if info is None: return None - return _decode_topk_buffer(info, len(sample.tokens) - 1, args.num_layers, args.moe_router_topk) + return _decode_topk_buffer(info, len(sample.tokens) - 1, args.num_layers, -1) def get_indexer_topk_from_response(args, output, sample): diff --git a/miles/rollout/sglang_rollout.py b/miles/rollout/sglang_rollout.py index fb3962a2aaf..2502e4e98b2 100644 --- a/miles/rollout/sglang_rollout.py +++ b/miles/rollout/sglang_rollout.py @@ -149,7 +149,10 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED ), f"Sample status is {sample.status}" - if state.processor and sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()): + if state.processor and ( + isinstance(sample.prompt, (list, tuple)) + or (sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values())) + ): processor_output = call_processor(state.processor, sample.prompt, sample.multimodal_inputs) prompt_ids = processor_output["input_ids"][0] prompt_ids = prompt_ids.tolist() if hasattr(prompt_ids, "tolist") else list(prompt_ids) @@ -204,6 +207,13 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A image_data = sample.multimodal_inputs["images"] payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data] + if sample.multimodal_inputs and sample.multimodal_inputs.get("audios"): + import base64 as _b64 + + payload["audio_data"] = [ + f"data:audio;base64,{_b64.b64encode(a).decode('ascii')}" for a in sample.multimodal_inputs["audios"] + ] + # Use existing tokens for multi-turn or tokenize the new prompt if len(sample.response) > 0: payload["input_ids"] = sample.tokens @@ -242,14 +252,22 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A sample.rollout_log_probs += new_response_log_probs if "routed_experts" in output["meta_info"]: - sample.rollout_routed_experts = np.frombuffer( + _re = np.frombuffer( pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")), dtype=np.int32, - ).reshape( - len(sample.tokens) - 1, - args.num_layers, - args.moe_router_topk, ) + _ntok = int(output["meta_info"]["prompt_tokens"]) + len(new_response_tokens) - 1 + _topk = _re.size // max(1, _ntok * args.num_layers) + if _re.size == (_ntok + 1) * args.num_layers * max(1, _topk): + # stop-edge: sglang also forwarded the final token; its routing rows + # feed no training position - drop the tail position. + _re = _re[: _ntok * args.num_layers * _topk] + assert _re.size == _ntok * args.num_layers * _topk, ( + f"routed_experts buffer {_re.size} != ntok({_ntok}) x layers({args.num_layers}) x topk({_topk}); " + f"prompt_tokens={output['meta_info'].get('prompt_tokens')} response={len(new_response_tokens)} " + f"unexpanded_tokens={len(sample.tokens)}" + ) + sample.rollout_routed_experts = _re.reshape(_ntok, args.num_layers, _topk) if "indexer_topk" in output["meta_info"]: sample.rollout_indexer_topk = get_indexer_topk_from_response(args, output, sample) diff --git a/miles/utils/chat_template_utils/inkling.py b/miles/utils/chat_template_utils/inkling.py new file mode 100644 index 00000000000..8a5ff0c47d0 --- /dev/null +++ b/miles/utils/chat_template_utils/inkling.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +import functools +import json +import os +from typing import Any + +_MODEL_TYPES = ("inkling_mm_model", "inkling_mm_model") + + +@functools.cache +def _read_model_type(name_or_path: str) -> str: + if not name_or_path: + return "" + config_path = os.path.join(name_or_path, "config.json") + if not os.path.isfile(config_path): + return "" + try: + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + except (OSError, UnicodeDecodeError, json.JSONDecodeError): + return "" + if not isinstance(config, dict): + return "" + return config.get("model_type", "") or "" + + +def is_inkling(tokenizer: Any) -> bool: + return _read_model_type(getattr(tokenizer, "name_or_path", "")) in _MODEL_TYPES + + +def is_inkling_checkpoint(name_or_path: str) -> bool: + return _read_model_type(name_or_path) in _MODEL_TYPES + + +@functools.cache +def fixed_chat_template() -> str: + """The bundled Inkling chat template (the checkpoint's, with the upstream fixes).""" + # imported lazily: tito_tokenizer imports template.py, which imports this module + from miles.utils.chat_template_utils.tito_tokenizer import TEMPLATE_DIR + + return (TEMPLATE_DIR / "inkling_fixed.jinja").read_text(encoding="utf-8") diff --git a/miles/utils/chat_template_utils/template.py b/miles/utils/chat_template_utils/template.py index 09f63ecf3e6..18dd7bdf06e 100644 --- a/miles/utils/chat_template_utils/template.py +++ b/miles/utils/chat_template_utils/template.py @@ -26,7 +26,7 @@ from sglang.srt.entrypoints.openai.protocol import Tool from transformers.utils.chat_template_utils import render_jinja_template -from miles.utils.chat_template_utils import deepseek +from miles.utils.chat_template_utils import deepseek, inkling def load_hf_chat_template(model_id: str) -> str: @@ -280,6 +280,19 @@ def apply_chat_template( **kwargs, ) + if inkling.is_inkling(tokenizer): + # the fixed template needs parsed tool-call arguments and handles the + # thinking-effort line, tool_calls, and the end-sampling token itself + return tokenizer.apply_chat_template( + normalize_tool_arguments(messages, "dict"), + chat_template=inkling.fixed_chat_template(), + tokenize=tokenize, + tools=extract_tool_dicts(tools), + return_dict=False, + add_generation_prompt=add_generation_prompt, + **kwargs, + ) + messages = normalize_tool_arguments(messages, "dict") tool_defs = extract_tool_dicts(tools) render_kwargs = dict(add_generation_prompt=add_generation_prompt, **kwargs) diff --git a/miles/utils/hf_config.py b/miles/utils/hf_config.py index cd081798fa2..130d75e4cb4 100644 --- a/miles/utils/hf_config.py +++ b/miles/utils/hf_config.py @@ -81,6 +81,11 @@ def register_hf_config_aliases() -> None: auto_cls.register(compat_config, compat_model_cls, exist_ok=alias.override_hf_native) _REGISTERED_ALIASES.add(alias.model_type) + try: + import sglang.srt.configs.inkling # noqa: F401 + except ImportError: + pass + def load_hf_config( checkpoint_path: str, diff --git a/miles/utils/processing_utils.py b/miles/utils/processing_utils.py index be0a9777b80..1f511b5b057 100644 --- a/miles/utils/processing_utils.py +++ b/miles/utils/processing_utils.py @@ -145,6 +145,13 @@ def extract_multimodal_train_inputs(processor_output: Mapping[str, Any]) -> dict def load_processor(name_or_path: str, **kwargs): + from miles.utils.chat_template_utils.inkling import is_inkling_checkpoint + + if is_inkling_checkpoint(name_or_path): + from miles_plugins.models.inkling.mm_processor import InklingTrainProcessor + + return InklingTrainProcessor(name_or_path) + try: proc = AutoProcessor.from_pretrained(name_or_path, **kwargs) except (OSError, ValueError) as e: @@ -159,6 +166,9 @@ def load_processor(name_or_path: str, **kwargs): def process_vision_info(prompt, processor): + if hasattr(processor, "extract_media"): + return processor.extract_media(prompt) + # TODO: temporary solution, will write image utils for miles later from qwen_vl_utils import process_vision_info as qwen_process_vision_info diff --git a/miles/utils/types.py b/miles/utils/types.py index 0d89ee092e2..208e6eefd12 100644 --- a/miles/utils/types.py +++ b/miles/utils/types.py @@ -202,7 +202,14 @@ def validate(self): if self.rollout_routed_experts is not None: actual = len(self.rollout_routed_experts) expect = len(self.tokens) - 1 - assert actual == expect, f"rollout_routed_experts length ({actual}) != len(tokens) - 1 ({expect})" + mm = self.multimodal_train_inputs or {} + extra = sum( + int(c) - 1 for key in ("mm_vision_num_patches", "mm_audio_num_tokens") for c in list(mm.get(key) or []) + ) + assert actual in (expect, expect + extra), ( + f"rollout_routed_experts length ({actual}) != len(tokens) - 1 ({expect})" + f" or media-expanded ({expect + extra})" + ) if self.rollout_indexer_topk is not None: actual = len(self.rollout_indexer_topk) expect = len(self.tokens) - 1 diff --git a/miles_plugins/mbridge/__init__.py b/miles_plugins/mbridge/__init__.py index 4dfd4ca796a..dcd80946736 100644 --- a/miles_plugins/mbridge/__init__.py +++ b/miles_plugins/mbridge/__init__.py @@ -3,6 +3,7 @@ from .glm4 import GLM4Bridge from .glm4moe import GLM4MoEBridge from .glm4moe_lite import GLM4MoELiteBridge +from .inkling import InklingBridge from .joyai_llm_flash import JoyAILLMFlashBridge from .mimo import MimoBridge from .qwen3_5 import Qwen3_5Bridge @@ -18,4 +19,5 @@ "DeepseekV32Bridge", "DeepseekV4Bridge", "JoyAILLMFlashBridge", + "InklingBridge", ] diff --git a/miles_plugins/mbridge/inkling.py b/miles_plugins/mbridge/inkling.py new file mode 100644 index 00000000000..02a8e027cea --- /dev/null +++ b/miles_plugins/mbridge/inkling.py @@ -0,0 +1,231 @@ +import re + +import torch + +from mbridge.core import register_model +from mbridge.models import Qwen2MoEBridge + + +def _deinterleave_w13(w13: torch.Tensor) -> torch.Tensor: + gate = w13[0::2] + up = w13[1::2] + return torch.cat([gate, up], dim=0) + + +@register_model("inkling_mm_model") +class InklingBridge(Qwen2MoEBridge): + _DIRECT_MAPPING = { + "embedding.word_embeddings.weight": "model.llm.embed.weight", + "embedding.embed_norm.weight": "model.llm.embed_norm.weight", + "decoder.final_layernorm.weight": "model.llm.norm.weight", + "output_layer.weight": "model.llm.unembed.weight", + } + + _ATTENTION_MAPPING = { + "self_attention.linear_qkv.layer_norm_weight": ["model.llm.layers.{layer_number}.attn_norm.weight"], + "self_attention.linear_qkv.weight": [ + "model.llm.layers.{layer_number}.attn.wq_du.weight", + "model.llm.layers.{layer_number}.attn.wk_dv.weight", + "model.llm.layers.{layer_number}.attn.wv_dv.weight", + "model.llm.layers.{layer_number}.attn.wr_du.weight", + ], + "self_attention.linear_proj.weight": ["model.llm.layers.{layer_number}.attn.wo_ud.weight"], + "self_attention.q_norm.weight": ["model.llm.layers.{layer_number}.attn.q_norm.weight"], + "self_attention.k_norm.weight": ["model.llm.layers.{layer_number}.attn.k_norm.weight"], + "self_attention.rel_proj": ["model.llm.layers.{layer_number}.attn.rel_logits_proj.proj"], + "self_attention.k_sconv.weight": ["model.llm.layers.{layer_number}.attn.k_sconv.weight"], + "self_attention.v_sconv.weight": ["model.llm.layers.{layer_number}.attn.v_sconv.weight"], + "self_attention.attn_sconv.weight": ["model.llm.layers.{layer_number}.attn_sconv.weight"], + } + + _MLP_MAPPING = { + "pre_mlp_layernorm.weight": ["model.llm.layers.{layer_number}.mlp_norm.weight"], + "mlp.mlp_sconv.weight": ["model.llm.layers.{layer_number}.mlp_sconv.weight"], + "mlp.router.weight": ["model.llm.layers.{layer_number}.mlp.gate.weight"], + "mlp.router.shared_gate": ["model.llm.layers.{layer_number}.mlp.gate.weight"], + "mlp.router.global_scale": ["model.llm.layers.{layer_number}.mlp.gate.global_scale"], + "mlp.router.expert_bias": ["model.llm.layers.{layer_number}.mlp.gate.bias"], + "mlp.experts.linear_fc1": ["model.llm.layers.{layer_number}.mlp.experts.w13_weight"], + "mlp.experts.linear_fc2": ["model.llm.layers.{layer_number}.mlp.experts.w2_weight"], + "mlp.linear_fc1.layer_norm_weight": ["model.llm.layers.{layer_number}.mlp_norm.weight"], + "mlp.linear_fc1.weight": ["model.llm.layers.{layer_number}.mlp.w13_dn.weight"], + "mlp.linear_fc2.weight": ["model.llm.layers.{layer_number}.mlp.w2_md.weight"], + "mlp.global_scale": ["model.llm.layers.{layer_number}.mlp.global_scale"], + "mlp.shared_experts.experts": [ + "model.llm.layers.{layer_number}.mlp.shared_experts.shared_w13_weight", + "model.llm.layers.{layer_number}.mlp.shared_experts.shared_w2_weight", + ], + } + + _CONFIG_MAPPING = { + "num_layers": "num_hidden_layers", + "hidden_size": "hidden_size", + "num_attention_heads": "num_attention_heads", + "num_query_groups": "num_key_value_heads", + "ffn_hidden_size": "intermediate_size", + "attention_dropout": ("attention_dropout", 0.0), + "layernorm_epsilon": "rms_norm_eps", + "hidden_dropout": ("hidden_dropout", 0.0), + "kv_channels": ("head_dim", None), + } + + def _get_text_config(self): + if hasattr(self.hf_config, "text_config"): + return self.hf_config.text_config + return self.hf_config + + def _is_tied_word_embeddings(self) -> bool: + return False + + def _adjust_mapping_for_shared_weights(self): + pass + + def _get_hf_shared_weight_keys(self): + return [] + + def _weight_name_mapping_mlp(self, name: str) -> list[str]: + layer_number = name.split(".")[2] + m = re.search(r"mlp\.shared_experts\.experts\.(\d+)\.(linear_fc1|linear_fc2)\.weight", name) + if m is not None: + which = m.group(2) + key = ( + "model.llm.layers.{layer_number}.mlp.shared_experts.shared_w13_weight" + if which == "linear_fc1" + else "model.llm.layers.{layer_number}.mlp.shared_experts.shared_w2_weight" + ) + return [key.format(layer_number=layer_number)] + convert_names = [] + for keyword, mapping_names in self._MLP_MAPPING.items(): + if keyword in name: + convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names]) + break + if len(convert_names) == 0: + raise NotImplementedError(f"Unsupported parameter name: {name}") + return convert_names + + def _weight_to_mcore_format(self, mcore_weights_name: str, hf_weights: list[torch.Tensor]) -> torch.Tensor: + if "self_attention.linear_qkv." in mcore_weights_name and "layer_norm" not in mcore_weights_name: + assert len(hf_weights) == 4, f"qkvr expects 4 HF tensors, got {len(hf_weights)}" + if getattr(self, "dtype", None) is not None: + hf_weights = [w.to(self.dtype) if w.dtype != self.dtype else w for w in hf_weights] + return torch.cat(hf_weights, dim=0).contiguous() + + if mcore_weights_name.endswith("mlp.router.weight"): + assert len(hf_weights) == 1 + tc = self._get_text_config() + nr = tc.n_routed_experts + w = hf_weights[0] + if getattr(self, "dtype", None) is not None and w.dtype != self.dtype: + w = w.to(self.dtype) + return w[:nr].contiguous() + if mcore_weights_name.endswith("mlp.router.shared_gate"): + assert len(hf_weights) == 1 + tc = self._get_text_config() + nr = tc.n_routed_experts + ns = tc.n_shared_experts + w = hf_weights[0] + if getattr(self, "dtype", None) is not None and w.dtype != self.dtype: + w = w.to(self.dtype) + return w[nr : nr + ns].contiguous() + if mcore_weights_name.endswith("global_scale"): + assert len(hf_weights) == 1 + return hf_weights[0].reshape(1).contiguous() + if mcore_weights_name.endswith("mlp.router.expert_bias"): + assert len(hf_weights) == 1 + return hf_weights[0].contiguous() + + def _global_eid(name): + nr = self._get_text_config().n_routed_experts + return self.mpu.ep_rank * (nr // self.mpu.ep_size) + int(name.split("weight")[-1]) + + if "mlp.experts.linear_fc1" in mcore_weights_name and len(hf_weights) == 1 and hf_weights[0].dim() == 3: + return _deinterleave_w13(hf_weights[0][_global_eid(mcore_weights_name)]).contiguous() + if "mlp.experts.linear_fc2" in mcore_weights_name and len(hf_weights) == 1 and hf_weights[0].dim() == 3: + return hf_weights[0][_global_eid(mcore_weights_name)].contiguous() + + sm = re.search(r"mlp\.shared_experts\.experts\.(\d+)\.(linear_fc1|linear_fc2)\.weight", mcore_weights_name) + if sm is not None and len(hf_weights) == 1 and hf_weights[0].dim() == 3: + e = int(sm.group(1)) + which = sm.group(2) + w = hf_weights[0][e] + if which == "linear_fc1": + w = _deinterleave_w13(w) + return w.contiguous() + + if mcore_weights_name.endswith("mlp.linear_fc1.weight") and len(hf_weights) == 1 and hf_weights[0].dim() == 2: + w = hf_weights[0] + if getattr(self, "dtype", None) is not None and w.dtype != self.dtype: + w = w.to(self.dtype) + return _deinterleave_w13(w).contiguous() + if mcore_weights_name.endswith("mlp.linear_fc2.weight") and len(hf_weights) == 1 and hf_weights[0].dim() == 2: + w = hf_weights[0] + if getattr(self, "dtype", None) is not None and w.dtype != self.dtype: + w = w.to(self.dtype) + return w.contiguous() + + return super()._weight_to_mcore_format(mcore_weights_name, hf_weights) + + def _weight_split_across_tp( + self, + mcore_weights_name: str, + mcore_weights: torch.Tensor, + param: torch.Tensor, + tp_split_size: int, + ) -> list[torch.Tensor]: + """Split fused linear_qkv into per-rank [q_i|k_i|v_i|r_i] blocks (base chunk(tp) scrambles them).""" + if ( + tp_split_size > 1 + and "self_attention.linear_qkv." in mcore_weights_name + and "layer_norm" not in mcore_weights_name + ): + tc = self._get_text_config() + nh = int(tc.num_attention_heads) + hd = int(tc.head_dim) + d_rel = int(getattr(tc, "d_rel", 16) or 16) + q_rows = nh * hd + r_rows = nh * d_rel + total = mcore_weights.shape[0] + kv_rows = (total - q_rows - r_rows) // 2 + nkv = kv_rows // hd if hd else 0 + assert ( + kv_rows > 0 + and (total - q_rows - r_rows) % 2 == 0 + and nkv * hd == kv_rows + and q_rows + 2 * kv_rows + r_rows == total + ), ( + f"qkv split {mcore_weights_name}: total={total} nh={nh} hd={hd} d_rel={d_rel} " + f"q={q_rows} kv={kv_rows} r={r_rows} nkv={nkv}" + ) + assert ( + nh % tp_split_size == 0 and nkv % tp_split_size == 0 + ), f"qkv tp={tp_split_size} must divide nh={nh} and nkv={nkv} for {mcore_weights_name}" + q, k, v, r = mcore_weights.split([q_rows, kv_rows, kv_rows, r_rows], dim=0) + qs, ks, vs, rs = ( + q.chunk(tp_split_size), + k.chunk(tp_split_size), + v.chunk(tp_split_size), + r.chunk(tp_split_size), + ) + return [torch.cat([qs[i], ks[i], vs[i], rs[i]], dim=0).contiguous() for i in range(tp_split_size)] + return super()._weight_split_across_tp(mcore_weights_name, mcore_weights, param, tp_split_size) + + def _build_config(self): + tc = self._get_text_config() + return self._build_base_config( + text_config_key="text_config" if hasattr(self.hf_config, "text_config") else None, + use_cpu_initialization=False, + num_moe_experts=tc.n_routed_experts, + moe_router_topk=tc.num_experts_per_tok, + moe_ffn_hidden_size=tc.intermediate_size, + moe_shared_expert_intermediate_size=tc.intermediate_size, + moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, + moe_router_pre_softmax=True, + moe_router_load_balancing_type="seq_aux_loss", + moe_aux_loss_coeff=0.0, + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + qk_layernorm=True, + add_qkv_bias=False, + add_bias_linear=False, + ) diff --git a/miles_plugins/models/inkling/__init__.py b/miles_plugins/models/inkling/__init__.py new file mode 100644 index 00000000000..18bde2f0e92 --- /dev/null +++ b/miles_plugins/models/inkling/__init__.py @@ -0,0 +1,19 @@ +def get_inkling_spec(*args, **kwargs): + from .model import get_inkling_spec as get_spec + + return get_spec(*args, **kwargs) + + +def inkling_model_provider(*args, **kwargs): + from .model import inkling_model_provider as provider + + return provider(*args, **kwargs) + + +def inkling_mm_model_provider(*args, **kwargs): + from .model import inkling_mm_model_provider as provider + + return provider(*args, **kwargs) + + +__all__ = ["get_inkling_spec", "inkling_mm_model_provider", "inkling_model_provider"] diff --git a/miles_plugins/models/inkling/audio_processing.py b/miles_plugins/models/inkling/audio_processing.py new file mode 100644 index 00000000000..0d756ad036f --- /dev/null +++ b/miles_plugins/models/inkling/audio_processing.py @@ -0,0 +1,181 @@ +from __future__ import annotations + +import io +import math +from collections.abc import Sequence +from dataclasses import dataclass + +import numpy as np +import torch +import torch.nn.functional as F + + +@dataclass +class InklingAudioEncoderParams: + sample_rate: int = 16_000 + window_size_multiplier: float = 2.0 + n_fft: int | None = None + n_mels: int = 80 + num_dmel_bins: int = 16 + dmel_min_value: float = -7.0 + dmel_max_value: float = 2.0 + audio_token_duration_s: float = 0.05 + + +def _load_audio_bytes(audio) -> bytes: + if isinstance(audio, (bytes, bytearray)): + return bytes(audio) + if hasattr(audio, "read"): + return audio.read() + if isinstance(audio, str): + path = audio[len("file://") :] if audio.startswith("file://") else audio + with open(path, "rb") as f: + return f.read() + raise TypeError(f"Unsupported audio input type for Inkling audio extractor: {type(audio)}") + + +def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int: + rounded = round(value) + if abs(value - rounded) > tolerance: + raise ValueError(f"{name} must resolve to an integer sample count, got {value}") + return int(rounded) + + +def _decode_audio(audio_bytes: bytes, sample_rate: int) -> torch.Tensor: + import soundfile as sf + + samples, src_sample_rate = sf.read(io.BytesIO(audio_bytes), dtype="float32", always_2d=True) + mono = samples.mean(axis=1) + if src_sample_rate != sample_rate: + mono = _resample(mono, src_sample_rate, sample_rate) + return torch.from_numpy(np.ascontiguousarray(mono, dtype=np.float32)) + + +def _resample(samples: np.ndarray, src_sample_rate: int, sample_rate: int) -> np.ndarray: + import torchaudio.functional as AF + + audio = torch.from_numpy(np.ascontiguousarray(samples, dtype=np.float32)) + resampled = AF.resample(audio, orig_freq=src_sample_rate, new_freq=sample_rate) + return resampled.detach().cpu().numpy().astype(np.float32, copy=False) + + +def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray: + """Slaney mel scale, matching the librosa/torchaudio convention.""" + frequencies = np.asarray(frequencies, dtype=np.float64) + f_sp = 200.0 / 3.0 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = np.log(6.4) / 27.0 + linear = frequencies / f_sp + log = min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep + return np.where(frequencies >= min_log_hz, log, linear) + + +def _mel_to_hz(mels: np.ndarray) -> np.ndarray: + mels = np.asarray(mels, dtype=np.float64) + f_sp = 200.0 / 3.0 + min_log_hz = 1000.0 + min_log_mel = min_log_hz / f_sp + logstep = np.log(6.4) / 27.0 + linear = mels * f_sp + log = min_log_hz * np.exp(logstep * (mels - min_log_mel)) + return np.where(mels >= min_log_mel, log, linear) + + +_MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {} + + +def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor: + key = (sample_rate, n_fft, n_mels) + cached = _MEL_BASIS_CACHE.get(key) + if cached is not None: + return cached + + fft_bins = n_fft // 2 + 1 + fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft + mel_edges = _mel_to_hz( + np.linspace( + _hz_to_mel(np.array([0.0]))[0], + _hz_to_mel(np.array([sample_rate / 2.0]))[0], + n_mels + 2, + dtype=np.float64, + ) + ) + mel_widths = np.diff(mel_edges) + lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None] + upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None] + weights = np.maximum(0.0, np.minimum(lower, upper)) + + weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None] + basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous() + _MEL_BASIS_CACHE[key] = basis + return basis + + +def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor: + hop_length = _to_exact_int( + params.audio_token_duration_s * params.sample_rate, + "audio_token_duration_s * sample_rate", + ) + window_size = _to_exact_int( + params.audio_token_duration_s * params.window_size_multiplier * params.sample_rate, + "audio_token_duration_s * window_size_multiplier * sample_rate", + ) + n_fft = params.n_fft or window_size + if hop_length <= 0 or window_size <= 0 or n_fft <= 0: + raise ValueError("audio hop length, window size, and n_fft must be positive") + if audio.numel() == 0: + return torch.empty((0, params.n_mels), dtype=torch.int32) + + right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel() + left_pad = max(n_fft - hop_length, 0) + audio = F.pad(audio, (left_pad, right_pad)) + + window = torch.hann_window(window_size, periodic=True, dtype=torch.float32) + spec = torch.stft( + audio.unsqueeze(0), + n_fft=n_fft, + hop_length=hop_length, + win_length=window_size, + window=window, + center=False, + normalized=False, + onesided=True, + return_complex=True, + ) + spec_ri = torch.view_as_real(spec) + magnitude = (spec_ri[..., 0].square() + spec_ri[..., 1].square()).clamp_min(1e-10).sqrt().squeeze(0) + + mel = _mel_basis(params.sample_rate, n_fft, params.n_mels).matmul(magnitude).clamp_min(1e-10).log10() + mel = mel.to(torch.float64).clamp(min=params.dmel_min_value, max=params.dmel_max_value) + bin_centers = torch.linspace( + params.dmel_min_value, + params.dmel_max_value, + params.num_dmel_bins, + dtype=torch.float64, + ) + dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1) + return dmel_bins.to(torch.int32).T.contiguous() + + +class InklingAudioDmelExtractor: + """audios -> {dmel_bins: [T_i, n_mels] int32 per clip, num_audio_tokens} (no HF deps).""" + + def __init__(self, params: dict | None = None): + merged = InklingAudioEncoderParams() + if params: + for k, v in params.items(): + if hasattr(merged, k): + setattr(merged, k, v) + self.params = merged + + def extract(self, audios: Sequence) -> dict: + if not isinstance(audios, (list, tuple)): + audios = [audios] + dmel_bins = [ + _dmel_bins(_decode_audio(_load_audio_bytes(a), self.params.sample_rate), self.params) for a in audios + ] + return { + "dmel_bins": dmel_bins, + "num_audio_tokens": [int(b.shape[0]) for b in dmel_bins], + } diff --git a/miles_plugins/models/inkling/image_processing.py b/miles_plugins/models/inkling/image_processing.py new file mode 100644 index 00000000000..2fbe5d0d0fb --- /dev/null +++ b/miles_plugins/models/inkling/image_processing.py @@ -0,0 +1,191 @@ +from __future__ import annotations + +import io +import math + +import numpy as np +import torch + +IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) +IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32) +PAD_RAW_VALUE = np.float32(-1.0 / 255.0) +PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD + +RESCALE_IMAGE_FRAC = 2.0 +RESCALE_IMAGE_MAX_UPSCALED_LONG_EDGE = 2048 + + +def _scaled_image_dimensions( + width: int, + height: int, + rescale_image_frac: float | None, + rescale_image_max_upscaled_long_edge: int | None, +) -> tuple[int, int]: + if rescale_image_frac is None: + return width, height + long_edge = max(width, height) + if long_edge == 0: + return width, height + target_long_edge = float(long_edge) * rescale_image_frac + if rescale_image_max_upscaled_long_edge is not None: + effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge) + target_long_edge = min(target_long_edge, float(effective_cap)) + ratio = target_long_edge / float(long_edge) + if ratio == 1.0: + return width, height + + def scale(value: int) -> int: + return max(1, math.floor(float(value) * ratio + 0.5)) + + return scale(width), scale(height) + + +def _load_image_bytes(image) -> bytes: + """Coerce a single image input into raw PNG/JPEG bytes for preprocessing.""" + if isinstance(image, (bytes, bytearray, memoryview)): + return bytes(image) + if isinstance(image, str): + if image.startswith(("http://", "https://", "data:")): + raise ValueError( + "Inkling image patchifier received a URL/data: image; resolve it to " + "bytes upstream before preprocessing." + ) + path = image[len("file://") :] if image.startswith("file://") else image + with open(path, "rb") as f: + return f.read() + + from PIL import Image + + if isinstance(image, torch.Tensor): + arr = image.detach().cpu().numpy() + image = Image.fromarray(arr.astype("uint8") if arr.dtype != np.uint8 else arr) + elif not isinstance(image, Image.Image): + image = Image.fromarray(image) + if image.mode != "RGB": + image = image.convert("RGB") + buf = io.BytesIO() + image.save(buf, format="PNG") + return buf.getvalue() + + +def _fill_patches_numpy( + arr: np.ndarray, + patch_size: int, + patches: np.ndarray, + mean: np.ndarray, + std: np.ndarray, + pad_norm: np.ndarray, +) -> None: + """Same math as the serving-side numba kernel: normalize in-image pixels, + pad-fill out-of-image cells; float32 throughout (bit-equal op order: + raw*inv255 -> -mean -> /std per channel).""" + h, w = arr.shape[0], arr.shape[1] + nph = (h + patch_size - 1) // patch_size + npw = w // patch_size + 1 + inv255 = np.float32(1.0 / 255.0) + + padded_h, padded_w = nph * patch_size, npw * patch_size + canvas = np.empty((padded_h, padded_w, 3), dtype=np.float32) + canvas[:] = pad_norm + norm = (arr.astype(np.float32) * inv255 - mean) / std + canvas[:h, : min(w, padded_w)] = norm[:, : min(w, padded_w)] + + view = canvas.reshape(nph, patch_size, npw, patch_size, 3).transpose(0, 2, 1, 3, 4) + patches[:] = view.reshape(nph * npw, patch_size, patch_size, 3) + + +try: + from numba import njit + + @njit(cache=True) + def _fill_patches_numba(arr, patch_size, patches, mean, std, pad_norm): # pragma: no cover + h = arr.shape[0] + w = arr.shape[1] + nph = (h + patch_size - 1) // patch_size + npw = w // patch_size + 1 + inv255 = np.float32(1.0 / 255.0) + + for k in range(nph * npw): + i = k // npw + j = k - i * npw + y_base = i * patch_size + x_base = j * patch_size + + for y in range(patch_size): + iy = y_base + y + for x in range(patch_size): + ix = x_base + x + if iy < h and ix < w: + for c in range(3): + raw = np.float32(arr[iy, ix, c]) * inv255 + patches[k, y, x, c] = (raw - mean[c]) / std[c] + else: + for c in range(3): + patches[k, y, x, c] = pad_norm[c] + + _fill_patches = _fill_patches_numba +except ImportError: # pragma: no cover + _fill_patches = _fill_patches_numpy + + +def encode_image_bytes(image_bytes: bytes, patch_size: int) -> torch.Tensor: + if patch_size <= 0: + raise ValueError("patch_size must be greater than zero") + + from PIL import Image + + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + scaled_size = _scaled_image_dimensions( + *image.size, + rescale_image_frac=RESCALE_IMAGE_FRAC, + rescale_image_max_upscaled_long_edge=RESCALE_IMAGE_MAX_UPSCALED_LONG_EDGE, + ) + if scaled_size != image.size: + image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS) + arr = np.array(image, dtype=np.uint8, copy=True) + height, width, _ = arr.shape + + nph = (height + patch_size - 1) // patch_size + npw = width // patch_size + 1 + num_patches = nph * npw + + patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32) + _fill_patches(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM) + + return ( + torch.from_numpy(patches) + .to(torch.bfloat16) + .view(num_patches, 1, patch_size, patch_size, 3) + .expand(num_patches, 2, patch_size, patch_size, 3) + ) + + +class InklingImagePatchifier: + """images -> {vision_patches_bthwc, num_patches, num_tokens} (dict, no HF deps).""" + + def __init__(self, patch_size: int = 40): + self.patch_size = patch_size + + def preprocess(self, images: list | object) -> dict: + if not isinstance(images, (list, tuple)): + images = [images] + + per_image_patches: list[torch.Tensor] = [] + num_patches: list[int] = [] + for img in images: + vp = encode_image_bytes(_load_image_bytes(img), patch_size=self.patch_size) + per_image_patches.append(vp) + num_patches.append(int(vp.shape[0])) + + if len(per_image_patches) == 1: + vision_patches_bthwc = per_image_patches[0] + elif per_image_patches: + vision_patches_bthwc = torch.cat(per_image_patches, dim=0) + else: + vision_patches_bthwc = torch.empty(0) + + return { + "vision_patches_bthwc": vision_patches_bthwc, + "num_patches": num_patches, + "num_tokens": list(num_patches), + } diff --git a/miles_plugins/models/inkling/layers.py b/miles_plugins/models/inkling/layers.py new file mode 100644 index 00000000000..2ece5c50de0 --- /dev/null +++ b/miles_plugins/models/inkling/layers.py @@ -0,0 +1,562 @@ +"""Inkling transformer layers: attention (dual global/SWA + rel bias + sconv), router +(sigmoid + shared-expert joint renorm + global_scale), shared experts, MoE/dense MLP.""" + +from __future__ import annotations + +import megatron.core.parallel_state as ps +import torch +import torch.nn as nn +import torch.nn.functional as F +import transformer_engine.pytorch as te +from megatron.core.extensions.transformer_engine import TELayerNormColumnParallelLinear, TERowParallelLinear +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +from megatron.core.transformer.attention import SelfAttention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.mlp import MLP +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.moe.moe_layer import MoELayer +from megatron.core.transformer.moe.router import TopKRouter + +from miles_plugins.models.inkling.ops import cp_all_gather, cp_world +from miles_plugins.models.inkling.ops import create_block_mask as _create_block_mask +from miles_plugins.models.inkling.ops import flex_compiled as _flex +from miles_plugins.models.inkling.ops import ( + inkling_sconv_fp32_packed, + inkling_sp_residual_conv, + inkling_sum_fp32, + inkling_swiglu_fp32, + seqlens_from_packed, +) + +_cp_world = cp_world +_cp_all_gather = cp_all_gather +_sp_residual_conv = inkling_sp_residual_conv +_seqlens_from_packed = seqlens_from_packed + + +class _Sconv(nn.Module): + def __init__(self, channels, kernel, dtype, impl="triton", packed=False): + super().__init__() + self.weight = nn.Parameter(torch.zeros(channels, 1, kernel, dtype=dtype)) + self.k = kernel + self.impl = impl + self.packed = packed + + def forward(self, x, seqlens=None): + return inkling_sconv_fp32_packed(x, self.weight, seqlens, impl=self.impl, packed=self.packed) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + + return make_sharded_tensors_for_checkpoint({"weight": self.weight}, prefix, sharded_offsets=sharded_offsets) + + +class InklingSelfAttention(SelfAttention): + def __init__(self, config, submodules, layer_number, attn_mask_type=AttnMaskType.causal, **kw): + super().__init__(config, submodules, layer_number, attn_mask_type=attn_mask_type, **kw) + t = config.inkling + self.is_local = (layer_number - 1) in t.local_layer_ids + self.nh = t.swa_num_attention_heads if self.is_local else t.num_attention_heads + self.nkv = t.swa_num_key_value_heads if self.is_local else t.num_key_value_heads + self.hd = t.swa_head_dim if self.is_local else t.head_dim + self.d_rel = t.d_rel + self.rel_extent = t.sliding_window_size if self.is_local else t.rel_extent + self.attn_backend = t.attn_backend + self._bm_cache = {} + tp = ps.get_tensor_model_parallel_world_size() + self.nh_l = self.nh // tp + self.nkv_l = max(1, self.nkv // tp) + qkvr_out = self.nh * self.hd + 2 * self.nkv * self.hd + self.nh * self.d_rel + self.linear_qkv = TELayerNormColumnParallelLinear( + config.hidden_size, + qkvr_out, + config=config, + init_method=config.init_method, + gather_output=False, + bias=False, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="qkv", + ) + self.linear_proj = TERowParallelLinear( + self.nh * self.hd, + config.hidden_size, + config=config, + init_method=config.output_layer_init_method, + bias=False, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="proj", + ) + self.q_norm = te.RMSNorm(self.hd, eps=t.rms_norm_eps, params_dtype=config.params_dtype) + self.k_norm = te.RMSNorm(self.hd, eps=t.rms_norm_eps, params_dtype=config.params_dtype) + self.rel_proj = nn.Parameter( + torch.zeros(self.d_rel, self.rel_extent, dtype=config.params_dtype), requires_grad=False + ) + self.window = (t.sliding_window_size - 1, 0) if self.is_local else (-1, -1) + self.dpa = te.DotProductAttention( + num_attention_heads=self.nh, + kv_channels=self.hd, + num_gqa_groups=self.nkv, + attention_dropout=0.0, + qkv_format="sbhd", + softmax_scale=1.0 / self.hd, + tp_size=tp, + tp_group=ps.get_tensor_model_parallel_group(), + ) + if t.use_sconv: + dt = config.params_dtype + sk = t.sconv_kernel_size + sconv_kw = {"impl": t.sconv_impl, "packed": t.sconv_packed} + self.k_sconv = _Sconv(self.nkv_l * self.hd, sk, dt, **sconv_kw) + self.v_sconv = _Sconv(self.nkv_l * self.hd, sk, dt, **sconv_kw) + for _w in (self.k_sconv.weight, self.v_sconv.weight): + _w.tensor_model_parallel = True + _w.partition_dim = 0 + _w.partition_stride = 1 + self.attn_sconv = _Sconv(config.hidden_size, sk, dt, **sconv_kw) + else: + self.k_sconv = self.v_sconv = self.attn_sconv = None + + def forward( + self, + hidden_states, + attention_mask, + inference_context=None, + rotary_pos_emb=None, + attention_bias=None, + packed_seq_params=None, + **kw, + ): + if hidden_states.dtype != self.config.params_dtype: + hidden_states = hidden_states.to(self.config.params_dtype) + sq = hidden_states.shape[0] + x = hidden_states.reshape(sq, -1) + qkvr, _ = self.linear_qkv(x) + T = qkvr.shape[0] + cp, cp_rank, cp_group = _cp_world() + q, k, v, r = qkvr.split( + [self.nh_l * self.hd, self.nkv_l * self.hd, self.nkv_l * self.hd, self.nh_l * self.d_rel], dim=-1 + ) + if cp > 1: + k = _cp_all_gather(k, cp_group, cp) + v = _cp_all_gather(v, cp_group, cp) + T_full = k.shape[0] + seqlens = _seqlens_from_packed(packed_seq_params, T_full) + self.config.inkling._seqlens = seqlens + if self.k_sconv is not None: + k = self.k_sconv(k, seqlens) + v = self.v_sconv(v, seqlens) + q = self.q_norm(q.reshape(-1, self.hd)).reshape(T, self.nh_l, self.hd) + k = self.k_norm(k.reshape(-1, self.hd)).reshape(T_full, self.nkv_l, self.hd) + v = v.reshape(T_full, self.nkv_l, self.hd) + r = r.reshape(T, self.nh_l, self.d_rel) + if self.attn_backend == "flex": + out = self._seg_cp_flex(q, k, v, r, seqlens, cp_rank * T, T_full) + else: + out = self._seg_cp(q, k, v, r, seqlens, cp_rank * T, T_full) + else: + seqlens = _seqlens_from_packed(packed_seq_params, T) + self.config.inkling._seqlens = seqlens + if self.k_sconv is not None: + k = self.k_sconv(k, seqlens) + v = self.v_sconv(v, seqlens) + q = self.q_norm(q.reshape(-1, self.hd)).reshape(T, self.nh_l, self.hd) + k = self.k_norm(k.reshape(-1, self.hd)).reshape(T, self.nkv_l, self.hd) + v = v.reshape(T, self.nkv_l, self.hd) + r = r.reshape(T, self.nh_l, self.d_rel) + out = self._attend(q, k, v, r, seqlens) + proj_out, proj_bias = self.linear_proj(out.unsqueeze(1)) + if self.attn_sconv is not None: + proj_out = _sp_residual_conv(self.config, self.attn_sconv, proj_out, seqlens) + return proj_out, proj_bias + + def _attend(self, q, k, v, r, seqlens): + if self.attn_backend == "fa4": + from miles_plugins.models.inkling.ops import inkling_fa4_attention as _fa4 + + T, nh_l, hd = q.shape + rel_logits = torch.einsum("thd,de->the", r.float(), self.rel_proj.float()) + out = _fa4(q, k, v, rel_logits, seqlens, self.window[0], self.is_local, 1.0 / self.hd, self.rel_extent) + return out.reshape(T, nh_l * hd) + if self.attn_backend == "flex": + return self._seg_flex(q, k, v, r, seqlens) + return self._seg(q, k, v, r, seqlens) + + def _seg(self, q, k, v, r, seqlens=None): + T, nh_l, hd = q.shape + rl = torch.einsum("thd,de->the", r.float(), self.rel_proj.float()) + RE = rl.shape[-1] + qi = torch.arange(T, device=q.device).view(T, 1) + ki = torch.arange(T, device=q.device).view(1, T) + rd = qi - ki + valid = (rd >= 0) & (rd < RE) + idx = rd.clamp(0, RE - 1) + rb = rl.permute(1, 0, 2) + bias = (torch.gather(rb, 2, idx.unsqueeze(0).expand(nh_l, T, T)) * valid.unsqueeze(0)).unsqueeze(0) + if seqlens is not None and len(seqlens) > 1: + seg_id = torch.repeat_interleave( + torch.arange(len(seqlens), device=q.device), torch.tensor(seqlens, device=q.device) + ) + cross = (seg_id.view(T, 1) != seg_id.view(1, T)).view(1, 1, T, T) + bias = bias.masked_fill(cross, -1e9) + ctx = self.dpa( + q.unsqueeze(1).contiguous(), + k.unsqueeze(1).contiguous(), + v.unsqueeze(1).contiguous(), + attn_mask_type="causal", + window_size=self.window, + core_attention_bias_type="post_scale_bias", + core_attention_bias=bias.to(q.dtype), + ) + return ctx.reshape(T, nh_l * hd) + + def _seg_cp(self, q, k, v, r, seqlens, q_off, T_full): + T_loc, nh_l, hd = q.shape + rl = torch.einsum("thd,de->the", r.float(), self.rel_proj.float()) + RE = rl.shape[-1] + qi = (q_off + torch.arange(T_loc, device=q.device)).view(T_loc, 1) + ki = torch.arange(T_full, device=q.device).view(1, T_full) + rd = qi - ki + idx = rd.clamp(0, RE - 1) + rb = rl.permute(1, 0, 2) + bias = torch.gather(rb, 2, idx.unsqueeze(0).expand(nh_l, T_loc, T_full)) * ((rd >= 0) & (rd < RE)).unsqueeze(0) + mask = rd < 0 + if self.is_local: + mask = mask | (rd > self.window[0]) + if seqlens is not None and len(seqlens) > 1: + seg = torch.repeat_interleave( + torch.arange(len(seqlens), device=q.device), torch.tensor(seqlens, device=q.device) + ) + mask = mask | (seg[q_off : q_off + T_loc].view(T_loc, 1) != seg.view(1, T_full)) + bias = bias.masked_fill(mask.unsqueeze(0), -1e9).unsqueeze(0).to(q.dtype) + ctx = self.dpa( + q.unsqueeze(1).contiguous(), + k.unsqueeze(1).contiguous(), + v.unsqueeze(1).contiguous(), + attn_mask_type="no_mask", + core_attention_bias_type="post_scale_bias", + core_attention_bias=bias, + ) + return ctx.reshape(T_loc, nh_l * hd) + + def _seg_id(self, seqlens, T, device): + if seqlens is None or len(seqlens) <= 1: + return None + return torch.repeat_interleave(torch.arange(len(seqlens), device=device), torch.tensor(seqlens, device=device)) + + def _block_mask(self, mask_mod, q_len, kv_len, key, device): + bm = self._bm_cache.get(key) + if bm is None: + if len(self._bm_cache) > 128: + self._bm_cache.clear() + bm = _create_block_mask( + mask_mod, + B=None, + H=None, + Q_LEN=q_len, + KV_LEN=kv_len, + device=device, + _compile=True, + ) + self._bm_cache[key] = bm + return bm + + def _seg_flex(self, q, k, v, r, seqlens=None): + T, nh_l, hd = q.shape + RE, W, is_local = self.rel_extent, self.window[0], self.is_local + rl = torch.einsum("thd,de->the", r.float(), self.rel_proj.float()) + rel_logits = rl.permute(1, 0, 2).contiguous() + seg_id = self._seg_id(seqlens, T, q.device) + + def score_mod(score, b, h, q_idx, kv_idx): + rd = q_idx - kv_idx + idx = torch.clamp(rd, 0, RE - 1) + bias = rel_logits[h, q_idx, idx] + valid = (rd >= 0) & (rd < RE) + return torch.where(valid, score + bias.to(score.dtype), score) + + def mask_mod(b, h, q_idx, kv_idx): + m = q_idx >= kv_idx + if seg_id is not None: + m = m & (seg_id[q_idx] == seg_id[kv_idx]) + if is_local: + m = m & (q_idx - kv_idx <= W) + return m + + block_mask = self._block_mask(mask_mod, T, T, (T, tuple(seqlens) if seqlens else None, is_local), q.device) + qf = q.permute(1, 0, 2).unsqueeze(0) + kf = k.permute(1, 0, 2).unsqueeze(0) + vf = v.permute(1, 0, 2).unsqueeze(0) + out = _flex()( + qf, + kf, + vf, + score_mod=score_mod, + block_mask=block_mask, + scale=1.0 / self.hd, + enable_gqa=True, + ) + return out.squeeze(0).permute(1, 0, 2).reshape(T, nh_l * hd) + + def _seg_cp_flex(self, q, k, v, r, seqlens, q_off, T_full): + T_loc, nh_l, hd = q.shape + RE, W, is_local = self.rel_extent, self.window[0], self.is_local + rl = torch.einsum("thd,de->the", r.float(), self.rel_proj.float()) + rel_logits = rl.permute(1, 0, 2).contiguous() + seg = self._seg_id(seqlens, T_full, q.device) + + def score_mod(score, b, h, q_idx, kv_idx): + rd = (q_off + q_idx) - kv_idx + idx = torch.clamp(rd, 0, RE - 1) + bias = rel_logits[h, q_idx, idx] + valid = (rd >= 0) & (rd < RE) + return torch.where(valid, score + bias.to(score.dtype), score) + + def mask_mod(b, h, q_idx, kv_idx): + gq = q_off + q_idx + rd = gq - kv_idx + m = rd >= 0 + if seg is not None: + m = m & (seg[gq] == seg[kv_idx]) + if is_local: + m = m & (rd <= W) + return m + + block_mask = self._block_mask( + mask_mod, + T_loc, + T_full, + (T_loc, T_full, tuple(seqlens) if seqlens else None, is_local, q_off), + q.device, + ) + qf = q.permute(1, 0, 2).unsqueeze(0) + kf = k.permute(1, 0, 2).unsqueeze(0) + vf = v.permute(1, 0, 2).unsqueeze(0) + out = _flex()( + qf, + kf, + vf, + score_mod=score_mod, + block_mask=block_mask, + scale=1.0 / self.hd, + enable_gqa=True, + ) + return out.squeeze(0).permute(1, 0, 2).reshape(T_loc, nh_l * hd) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + from megatron.core.dist_checkpointing import ShardedTensor + from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory + from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint + + sd = super().sharded_state_dict(prefix, sharded_offsets, metadata) + tp_rank = ps.get_tensor_model_parallel_rank() + tp_size = ps.get_tensor_model_parallel_world_size() + qkey = f"{prefix}linear_qkv.weight" + if qkey in sd: + orig = sd[qkey] + splits = [self.nh_l * self.hd, self.nkv_l * self.hd, self.nkv_l * self.hd, self.nh_l * self.d_rel] + so = tuple(sharded_offsets) + prepend = len(so) + + def _qkvr_build( + key, t, replica_id, flattened_range, _s=splits, _so=so, _p=prepend, _tr=tp_rank, _ts=tp_size + ): + comps = torch.split(t, _s, dim=0) + return [ + ShardedTensor.from_rank_offsets( + f"{key}.{nm}", c.contiguous(), *_so, (_p, _tr, _ts), replica_id=replica_id, prepend_axis_num=_p + ) + for nm, c in zip(("q", "k", "v", "r"), comps, strict=False) + ] + + def _qkvr_merge(sub): + return torch.cat(list(sub), dim=0) + + sd[qkey] = ShardedTensorFactory( + orig.key, orig.data, _qkvr_build, _qkvr_merge, orig.replica_id, flattened_range=orig.flattened_range + ) + if self.k_sconv is not None: + for sub in ("k_sconv", "v_sconv"): + key = f"{prefix}{sub}.weight" + if key in sd: + sd[key] = make_tp_sharded_tensor_for_checkpoint( + getattr(self, sub).weight, key, tp_axis=0, prepend_offsets=sharded_offsets + ) + return sd + + +class InklingRouter(TopKRouter): + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + H = self.config.hidden_size + ns = self.config.inkling.n_shared_experts + self.shared_gate = nn.Parameter(torch.zeros(ns, H, dtype=self.config.params_dtype)) + _gs_dtype = torch.float32 + self.global_scale = nn.Parameter(torch.ones(1, dtype=_gs_dtype)) + if _gs_dtype == torch.float32: + self.global_scale._keep_fp32 = True + self._cache_key = None + self._cache = None + + def forward(self, input, padding_mask=None, input_ids=None): + key = (input.data_ptr(), tuple(input.shape)) + if self._cache_key == key and self._cache is not None: + probs, routing_map = self._cache + self._cache_key = None + self._cache = None + return probs, routing_map + + self._maintain_float32_expert_bias() + H = input.shape[-1] + nr, topk = self.config.num_moe_experts, self.topk + logits = self.gating(input).view(-1, nr).float() + shared_logits = input.reshape(-1, H).float() @ self.shared_gate.float().t() + score = logits.sigmoid() + self.expert_bias.float() + from miles.utils.replay_base import routing_replay_manager + + _sel_topk = routing_replay_manager.get_topk_fn(lambda s, k: s.topk(k, dim=-1).indices, return_probs=False) + topk_ids = _sel_topk(score, topk).long() + sel_logits = logits.gather(-1, topk_ids) + active = torch.cat([sel_logits, shared_logits], dim=-1) + lp = F.logsigmoid(active) + w = torch.softmax(lp, dim=-1) * float(self.config.inkling.route_scale) * self.global_scale.float() + routed_w, shared_w = w[:, :topk], w[:, topk:] + probs = torch.zeros_like(logits).scatter(-1, topk_ids, routed_w) + routing_map = torch.zeros_like(logits, dtype=torch.bool).scatter(-1, topk_ids, True) + self.config.inkling._shared_w = shared_w + self._apply_expert_bias(routing_map, padding_mask) + if self.config.moe_router_dtype in ("fp32", "fp64"): + return probs, routing_map + return probs.type_as(input), routing_map + + +class InklingSharedExperts(MegatronModule): + def __init__(self, config, submodules=None, gate=False, pg_collection=None): + super().__init__(config=config) + ns = config.inkling.n_shared_experts + inter = config.inkling.intermediate_size + tp = pg_collection.tp if pg_collection is not None else None + import copy as _copy + + shared_cfg = _copy.copy(config) + shared_cfg.ffn_hidden_size = inter + self.experts = nn.ModuleList( + [MLP(config=shared_cfg, submodules=submodules, ffn_hidden_size=inter, tp_group=tp) for _ in range(ns)] + ) + + def forward(self, hidden_states): + sw = self.config.inkling._shared_w + s, b, h = hidden_states.shape + sw_ = sw.view(s, b, -1) + if self.config.sequence_parallel and ps.get_tensor_model_parallel_world_size() > 1: + sw_ = gather_from_sequence_parallel_region(sw_) + sw32 = sw_.float() + ys = [] + for j, mlp in enumerate(self.experts): + fc1_out, _ = mlp.linear_fc1(hidden_states) + act = inkling_swiglu_fp32(fc1_out, per_token_scale=sw32[:, :, j : j + 1]) + yj, _ = mlp.linear_fc2(act) + ys.append(yj) + return inkling_sum_fp32(ys) + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + sd = super().sharded_state_dict(prefix, sharded_offsets, metadata) + from megatron.core.transformer.mlp import apply_swiglu_sharded_factory + from megatron.core.utils import make_tp_sharded_tensor_for_checkpoint + + singleton = (metadata or {}).get("singleton_local_shards", False) + for j in range(len(self.experts)): + k1 = f"{prefix}experts.{j}.linear_fc1.weight" + if k1 in sd: + base = make_tp_sharded_tensor_for_checkpoint( + self.experts[j].linear_fc1.weight, k1, tp_axis=0, prepend_offsets=sharded_offsets + ) + sd[k1] = apply_swiglu_sharded_factory(base, sharded_offsets, singleton) + k2 = f"{prefix}experts.{j}.linear_fc2.weight" + if k2 in sd: + sd[k2] = make_tp_sharded_tensor_for_checkpoint( + self.experts[j].linear_fc2.weight, k2, tp_axis=1, prepend_offsets=sharded_offsets + ) + return sd + + +class InklingMoELayer(MoELayer): + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + t = self.config.inkling + self.mlp_sconv = ( + _Sconv( + self.config.hidden_size, + t.sconv_kernel_size, + self.config.params_dtype, + impl=t.sconv_impl, + packed=t.sconv_packed, + ) + if t.use_sconv + else None + ) + + def shared_experts_compute(self, hidden_states): + if self.use_shared_expert and not self.shared_expert_overlap: + probs, routing_map = self.router(hidden_states) + self.router._cache_key = (hidden_states.data_ptr(), tuple(hidden_states.shape)) + self.router._cache = (probs, routing_map) + return super().shared_experts_compute(hidden_states) + + def forward(self, hidden_states, *args, **kw): + if hidden_states.dtype != self.config.params_dtype: + hidden_states = hidden_states.to(self.config.params_dtype) + out, bias = super().forward(hidden_states, *args, **kw) + if self.mlp_sconv is not None: + seqlens = getattr(self.config.inkling, "_seqlens", None) + out = _sp_residual_conv(self.config, self.mlp_sconv, out, seqlens) + return out, bias + + +class InklingDenseMLP(MLP): + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + t = self.config.inkling + self.mlp_sconv = ( + _Sconv( + self.config.hidden_size, + t.sconv_kernel_size, + self.config.params_dtype, + impl=t.sconv_impl, + packed=t.sconv_packed, + ) + if t.use_sconv + else None + ) + self.global_scale = ( + nn.Parameter(torch.ones(1, dtype=self.config.params_dtype)) + if getattr(t, "use_global_scale", False) + else None + ) + + def forward(self, hidden_states, *args, **kw): + if hidden_states.dtype != self.config.params_dtype: + hidden_states = hidden_states.to(self.config.params_dtype) + fc1_out, _ = self.linear_fc1(hidden_states) + act = inkling_swiglu_fp32(fc1_out) + out, bias = self.linear_fc2(act) + if self.global_scale is not None: + out = out * self.global_scale.to(out.dtype) + if self.mlp_sconv is not None: + seqlens = getattr(self.config.inkling, "_seqlens", None) + out = _sp_residual_conv(self.config, self.mlp_sconv, out, seqlens) + return out, bias + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None): + sd = super().sharded_state_dict(prefix, sharded_offsets, metadata) + if self.global_scale is not None: + from megatron.core.transformer.utils import make_sharded_tensors_for_checkpoint + + sd.update( + make_sharded_tensors_for_checkpoint( + {"global_scale": self.global_scale}, prefix, sharded_offsets=sharded_offsets + ) + ) + return sd diff --git a/miles_plugins/models/inkling/mm_processor.py b/miles_plugins/models/inkling/mm_processor.py new file mode 100644 index 00000000000..9719dd4708c --- /dev/null +++ b/miles_plugins/models/inkling/mm_processor.py @@ -0,0 +1,268 @@ +from __future__ import annotations + +import json +import os +from typing import Any + +import torch + +from miles_plugins.models.inkling.image_processing import InklingImagePatchifier + +IMAGE_TOKEN_ID = -101 +AUDIO_TOKEN_ID = -102 +ROLE_MESSAGE_TOKENS = { + "user": "<|message_user|>", + "assistant": "<|message_model|>", + "system": "<|message_system|>", + "tool": "<|message_tool|>", +} +MESSAGE_MODEL = "<|message_model|>" +CONTENT_TEXT = "<|content_text|>" +CONTENT_THINKING = "<|content_thinking|>" +CONTENT_XML = "<|content_xml|>" +CONTENT_IMAGE = "<|content_image|>" +CONTENT_AUDIO_INPUT = "<|content_audio_input|>" +END_MESSAGE = "<|end_message|>" +AUDIO_END = "<|audio_end|>" + +_IMAGE_PART_TYPES = frozenset({"image", "input_image", "image_url"}) +_AUDIO_PART_TYPES = frozenset({"audio", "input_audio", "audio_url"}) + + +class _TokenizerAdapter: + """encode_special/encode_text over the HF o200k+overlay tokenizer.""" + + def __init__(self, tokenizer): + self._tok = tokenizer + self._special_cache: dict[str, int] = {} + + def encode_special(self, literal: str) -> int: + tid = self._special_cache.get(literal) + if tid is None: + ids = self._tok.encode(literal, add_special_tokens=False) + assert len(ids) == 1, f"Inkling special {literal!r} must encode to one id, got {ids}" + tid = self._special_cache[literal] = ids[0] + return tid + + def encode_text(self, text: str) -> list[int]: + return self._tok.encode(text, add_special_tokens=False) + + +def _iter_render_parts(content: Any): + """Yield (kind, text) per content part: kind in {text, image, audio}.""" + if content is None: + return + if isinstance(content, str): + if content: + yield ("text", content) + return + if not isinstance(content, (list, tuple)): + raise TypeError("message content must be a string or a sequence of parts") + for part in content: + if isinstance(part, str): + yield ("text", part) + continue + if not isinstance(part, dict): + raise TypeError(f"content part must be a dict, got {type(part).__name__}") + ptype = part.get("type") + if ptype in (None, "text", "input_text"): + text = part.get("text", "") + yield ("text", text if isinstance(text, str) else "") + elif ptype in _IMAGE_PART_TYPES: + yield ("image", "") + elif ptype in _AUDIO_PART_TYPES: + yield ("audio", "") + else: + raise ValueError(f"unsupported content part type: {ptype!r}") + + +def _append_message( + input_ids: list[int], + tok: _TokenizerAdapter, + role: str, + kind: str, + text: str, + *, + author_name: str | None = None, +) -> None: + input_ids.append(tok.encode_special(ROLE_MESSAGE_TOKENS[role])) + if author_name: + input_ids.extend(tok.encode_text(author_name)) + + if kind == "text": + input_ids.append(tok.encode_special(CONTENT_TEXT)) + input_ids.extend(tok.encode_text(text)) + elif kind == "image": + input_ids.append(tok.encode_special(CONTENT_IMAGE)) + input_ids.append(IMAGE_TOKEN_ID) + elif kind == "audio": + input_ids.append(tok.encode_special(CONTENT_AUDIO_INPUT)) + input_ids.append(AUDIO_TOKEN_ID) + input_ids.append(tok.encode_special(AUDIO_END)) + elif kind == "thinking": + input_ids.append(tok.encode_special(CONTENT_THINKING)) + input_ids.extend(tok.encode_text(text)) + elif kind == "xml": + input_ids.append(tok.encode_special(CONTENT_XML)) + input_ids.extend(tok.encode_text(text)) + else: + raise ValueError(f"unsupported Inkling render part kind: {kind!r}") + + input_ids.append(tok.encode_special(END_MESSAGE)) + + +_EFFORT_LEVELS = {"none": 0.0, "low": 0.2, "medium": 0.7, "high": 0.9, "xhigh": 0.99, "max": 0.99} + + +def _reasoning_effort_from_env() -> float | None: + val = os.environ.get("INKLING_REASONING_EFFORT", "").strip().lower() + if not val: + return None + if val in _EFFORT_LEVELS: + return _EFFORT_LEVELS[val] + return max(0.0, min(0.99, float(val))) + + +def render_inkling_messages_to_ids( + messages: list[dict[str, Any]], + tok: _TokenizerAdapter, + *, + tools: list[dict] | None = None, + add_generation_prompt: bool = True, + reasoning_effort: float | None = None, +) -> list[int]: + """Token-level Inkling chat render with one media sentinel per image/audio part.""" + input_ids: list[int] = [] + + if tools: + _append_message( + input_ids, + tok, + "system", + "xml", + json.dumps(tools, ensure_ascii=False), + author_name="tool_declare", + ) + + if reasoning_effort is not None: + _append_message( + input_ids, + tok, + "system", + "text", + f"Thinking effort level: {reasoning_effort}", + ) + + for msg in messages: + role = msg.get("role") + if role not in ROLE_MESSAGE_TOKENS: + raise ValueError(f"unsupported Inkling message role {role!r}") + if role == "assistant" and msg.get("reasoning_content"): + rc = msg["reasoning_content"] + if not isinstance(rc, str): + raise TypeError("assistant reasoning_content must be a string") + _append_message(input_ids, tok, role, "thinking", rc) + for kind, text in _iter_render_parts(msg.get("content", "")): + _append_message(input_ids, tok, role, kind, text) + + if add_generation_prompt: + input_ids.append(tok.encode_special(MESSAGE_MODEL)) + return input_ids + + +def _load_pil_image(spec): + """Resolve one image content part's payload to a PIL image (no resize).""" + from PIL import Image + + if isinstance(spec, Image.Image): + return spec + if isinstance(spec, str): + path = spec[len("file://") :] if spec.startswith("file://") else spec + if path.startswith(("http://", "https://", "data:")): + raise ValueError(f"Inkling extract_media: resolve URL images upstream, got {spec[:64]!r}") + img = Image.open(path) + img.load() + return img + raise TypeError(f"unsupported image payload type {type(spec).__name__}") + + +class InklingTrainProcessor: + """Duck-typed processor for miles' generic multimodal pipeline (Inkling).""" + + def __init__(self, hf_checkpoint: str): + from miles.utils.processing_utils import load_tokenizer + + with open(os.path.join(hf_checkpoint, "config.json"), encoding="utf-8") as f: + cfg = json.load(f) + vision_cfg = cfg.get("vision_config") or {} + self.patchifier = InklingImagePatchifier(patch_size=vision_cfg.get("patch_size", 40)) + audio_cfg = cfg.get("audio_config") or {} + from miles_plugins.models.inkling.audio_processing import InklingAudioDmelExtractor + + self.dmel_extractor = InklingAudioDmelExtractor( + params={ + "n_mels": audio_cfg.get("n_mel_bins", 80), + "num_dmel_bins": audio_cfg.get("mel_vocab_size", 16), + **{k: audio_cfg[k] for k in ("dmel_min_value", "dmel_max_value") if k in audio_cfg}, + } + ) + self.tok = _TokenizerAdapter(load_tokenizer(hf_checkpoint, trust_remote_code=True)) + + def extract_media(self, prompt: list[dict[str, Any]]) -> dict: + images, audios = [], [] + for msg in prompt: + content = msg.get("content", "") + if not isinstance(content, (list, tuple)): + continue + for part in content: + if not isinstance(part, dict): + continue + ptype = part.get("type") + if ptype in _IMAGE_PART_TYPES: + spec = part.get("image", part.get("image_url")) + if isinstance(spec, dict): + spec = spec.get("url") + images.append(_load_pil_image(spec)) + elif ptype in _AUDIO_PART_TYPES: + spec = part.get("audio", part.get("audio_url")) + if isinstance(spec, dict): + spec = spec.get("url") + if isinstance(spec, str): + path = spec[len("file://") :] if spec.startswith("file://") else spec + with open(path, "rb") as f: + spec = f.read() + assert isinstance( + spec, (bytes, bytearray) + ), f"audio payload must be a path or bytes, got {type(spec).__name__}" + audios.append(bytes(spec)) + return {"images": images or None, "audios": audios or None, "videos": None} + + def __call__(self, text=None, images=None, videos=None, audios=None, **kwargs): + assert videos is None or not videos, "Inkling processor: video inputs unsupported" + assert isinstance(text, (list, tuple)), ( + "Inkling processor expects the raw message list as `text` (run without " + f"--apply-chat-template so the prompt stays structured), got {type(text).__name__}" + ) + input_ids = render_inkling_messages_to_ids(list(text), self.tok, reasoning_effort=_reasoning_effort_from_env()) + + out: dict[str, Any] = {"input_ids": [input_ids]} + n_img_sentinels = sum(1 for t in input_ids if t == IMAGE_TOKEN_ID) + n_images = len(images) if images else 0 + assert ( + n_img_sentinels == n_images + ), f"Inkling render produced {n_img_sentinels} image sentinel(s) but {n_images} image(s) given" + if n_images: + feat = self.patchifier.preprocess(list(images)) + out["mm_vision_patches"] = feat["vision_patches_bthwc"] + out["mm_vision_num_patches"] = torch.tensor(feat["num_patches"], dtype=torch.long) + + n_aud_sentinels = sum(1 for t in input_ids if t == AUDIO_TOKEN_ID) + n_audios = len(audios) if audios else 0 + assert ( + n_aud_sentinels == n_audios + ), f"Inkling render produced {n_aud_sentinels} audio sentinel(s) but {n_audios} audio(s) given" + if n_audios: + feat = self.dmel_extractor.extract(list(audios)) + out["mm_audio_dmel"] = torch.cat(feat["dmel_bins"], dim=0) + out["mm_audio_num_tokens"] = torch.tensor(feat["num_audio_tokens"], dtype=torch.long) + return out diff --git a/miles_plugins/models/inkling/mm_towers.py b/miles_plugins/models/inkling/mm_towers.py new file mode 100644 index 00000000000..8229dd36bd3 --- /dev/null +++ b/miles_plugins/models/inkling/mm_towers.py @@ -0,0 +1,168 @@ +from __future__ import annotations + +import json + +import torch +import torch.nn as nn + +from miles_plugins.models.inkling.vision_encoder import HMLPPatchEncoder, InklingVisionConfig, RMSNorm + + +class InklingAudioEncoder(nn.Module): + def __init__(self, audio_cfg: dict): + super().__init__() + assert audio_cfg.get("audio_mode", "dmel") == "dmel" + self.n_mel_bins = audio_cfg["n_mel_bins"] + self.mel_vocab_size = audio_cfg["mel_vocab_size"] + self.encoder = nn.Embedding(self.n_mel_bins * self.mel_vocab_size, audio_cfg["decoder_dmodel"]) + self.final_norm = ( + RMSNorm(audio_cfg["decoder_dmodel"], eps=1e-6) if audio_cfg.get("use_audio_norm", True) else None + ) + + def forward(self, dmel: torch.Tensor) -> torch.Tensor: + assert dmel.shape[1] == self.n_mel_bins + idx = (torch.arange(self.n_mel_bins, device=dmel.device) * self.mel_vocab_size).unsqueeze(0) + dmel.to( + torch.int32 + ) + h = self.encoder(idx.reshape(-1)).reshape(dmel.shape[0], self.n_mel_bins, -1).sum(axis=1) + return self.final_norm(h) if self.final_norm is not None else h + + +class InklingVisionEncoder(nn.Module): + def __init__(self, vision_cfg: dict): + super().__init__() + assert vision_cfg.get("vision_encoder_type", "hmlp") == "hmlp" + self.vision_encoder = HMLPPatchEncoder(InklingVisionConfig(vision_cfg)) + + def forward(self, patches: torch.Tensor) -> torch.Tensor: + return self.vision_encoder(patches) + + +def _load_tower_tensors(hf_checkpoint: str) -> dict[str, torch.Tensor]: + """Read tower tensors from HF safetensors, normalizing keys to visual.* / audio.*.""" + from safetensors import safe_open + + index = json.load(open(f"{hf_checkpoint}/model.safetensors.index.json")) + by_file: dict[str, list[str]] = {} + for name, fname in index["weight_map"].items(): + bare = name[len("model.") :] if name.startswith(("model.visual.", "model.audio.")) else name + if bare.startswith(("visual.", "audio.")): + by_file.setdefault(fname, []).append(name) + out: dict[str, torch.Tensor] = {} + for fname, names in by_file.items(): + with safe_open(f"{hf_checkpoint}/{fname}", framework="pt") as f: + for name in names: + bare = name[len("model.") :] if name.startswith(("model.visual.", "model.audio.")) else name + out[bare] = f.get_tensor(name) + return out + + +def build_mm_towers( + hf_checkpoint: str, device, dtype, train: bool = False +) -> tuple[nn.Module | None, nn.Module | None]: + """Construct + HF-load both towers; None when the checkpoint has no config for a modality.""" + cfg = json.load(open(f"{hf_checkpoint}/config.json")) + weights = _load_tower_tensors(hf_checkpoint) + + visual = audio = None + if cfg.get("vision_config") and any(k.startswith("visual.") for k in weights): + visual = InklingVisionEncoder(cfg["vision_config"]) + sd = {k[len("visual.") :]: v for k, v in weights.items() if k.startswith("visual.")} + visual.vision_encoder.load_state_dict(sd, strict=True) + if cfg.get("audio_config") and any(k.startswith("audio.") for k in weights): + audio = InklingAudioEncoder(cfg["audio_config"]) + sd = {k[len("audio.") :]: v for k, v in weights.items() if k.startswith("audio.")} + audio.load_state_dict(sd, strict=True) + + for tower in (visual, audio): + if tower is not None: + tower.to(device=device, dtype=dtype) + tower.requires_grad_(train) + if not train: + tower.eval() + return visual, audio + + +def wire_mm_towers(model, hf_checkpoint: str, train: bool = False) -> None: + """Wrap an InklingGPTModel's forward for multimodal training (pre_process builds+merges; + later stages pass through). train=True trains the towers instead of freezing them + (weight-sync/ckpt for tower params is not wired yet).""" + import megatron.core.parallel_state as ps + + _orig_forward = model.forward + + if not getattr(model, "pre_process", False): + + def _mm_passthrough( + *a, + mm_vision_patches=None, + mm_vision_positions=None, + mm_vision_num_patches=None, + mm_audio_dmel=None, + mm_audio_positions=None, + mm_audio_num_tokens=None, + **kw, + ): + return _orig_forward(*a, **kw) + + model.forward = _mm_passthrough + return + + device = torch.cuda.current_device() + visual, audio = build_mm_towers(hf_checkpoint, device=device, dtype=model.config.params_dtype, train=train) + model.__dict__["_mm_towers"] = (visual, audio) + + def _scatter(decoder_input, embeds, positions): + if embeds is None or positions is None or positions.numel() == 0: + return decoder_input + s_local = decoder_input.shape[0] + if model.config.sequence_parallel and ps.get_tensor_model_parallel_world_size() > 1: + rank = ps.get_tensor_model_parallel_rank() + local = positions - rank * s_local + sel = (local >= 0) & (local < s_local) + local, emb = local[sel], embeds[sel] + else: + local, emb = positions, embeds + decoder_input = decoder_input.clone() + if local.numel(): + decoder_input[local, 0] = emb.to(decoder_input.dtype) + return decoder_input + + def _mm_forward( + *a, + mm_vision_patches=None, + mm_vision_positions=None, + mm_vision_num_patches=None, + mm_audio_dmel=None, + mm_audio_positions=None, + mm_audio_num_tokens=None, + **kw, + ): + input_ids = kw.get("input_ids", a[0] if a else None) + position_ids = kw.get("position_ids") + if mm_vision_patches is not None: + assert mm_vision_positions is not None, "mm_vision_patches given without mm_vision_positions" + assert mm_vision_positions.numel() == mm_vision_patches.shape[0], ( + f"{mm_vision_positions.numel()} vision position(s) vs " f"{mm_vision_patches.shape[0]} patch row(s)" + ) + if mm_audio_dmel is not None: + assert mm_audio_positions is not None, "mm_audio_dmel given without mm_audio_positions" + assert mm_audio_positions.numel() == mm_audio_dmel.shape[0], ( + f"{mm_audio_positions.numel()} audio position(s) vs " f"{mm_audio_dmel.shape[0]} dmel frame(s)" + ) + decoder_input = model.embedding(input_ids=input_ids, position_ids=position_ids) + # the grad context must cover ONLY the tower forwards: _scatter clones + # decoder_input, and a clone under no_grad detaches the embedding graph + ctx = torch.enable_grad() if train else torch.no_grad() + if mm_vision_patches is not None: + with ctx: + v_emb = model.__dict__["_mm_towers"][0](mm_vision_patches.to(device=decoder_input.device)) + decoder_input = _scatter(decoder_input, v_emb, mm_vision_positions.to(decoder_input.device)) + if mm_audio_dmel is not None: + with ctx: + a_emb = model.__dict__["_mm_towers"][1](mm_audio_dmel.to(device=decoder_input.device)) + decoder_input = _scatter(decoder_input, a_emb, mm_audio_positions.to(decoder_input.device)) + kw["decoder_input"] = decoder_input + return _orig_forward(*a, **kw) + + model.forward = _mm_forward diff --git a/miles_plugins/models/inkling/model.py b/miles_plugins/models/inkling/model.py new file mode 100644 index 00000000000..fe17932d0b3 --- /dev/null +++ b/miles_plugins/models/inkling/model.py @@ -0,0 +1,306 @@ +"""Inkling model assembly: TransformerConfig construction, layer/block specs, and the +GPTModel subclass + provider that miles' custom-model-provider path loads.""" + +from __future__ import annotations + +import logging +import os + +import torch +import transformer_engine.pytorch as te +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + +from miles_plugins.models.inkling.layers import ( + InklingDenseMLP, + InklingMoELayer, + InklingRouter, + InklingSelfAttention, + InklingSharedExperts, +) + +logger = logging.getLogger(__name__) + + +class InklingExtra: + def __init__(self, t: dict): + self.num_attention_heads = t["num_attention_heads"] + self.num_key_value_heads = t["num_key_value_heads"] + self.head_dim = t["head_dim"] + self.swa_num_attention_heads = t["swa_num_attention_heads"] + self.swa_num_key_value_heads = t["swa_num_key_value_heads"] + self.swa_head_dim = t["swa_head_dim"] + self.sliding_window_size = t["sliding_window_size"] + self.d_rel = t["d_rel"] + self.rel_extent = t.get("rel_extent", 1024) + self.local_layer_ids = set(t["local_layer_ids"]) + self.sconv_kernel_size = t["sconv_kernel_size"] + self.use_sconv = t["use_sconv"] + self.use_embed_norm = t["use_embed_norm"] + self.n_routed_experts = t["n_routed_experts"] + self.num_experts_per_tok = t["num_experts_per_tok"] + self.n_shared_experts = t["n_shared_experts"] + self.route_scale = t["route_scale"] + self.intermediate_size = t["intermediate_size"] + self.hidden_size = t["hidden_size"] + self.num_hidden_layers = t["num_hidden_layers"] + self.vocab_size = t["vocab_size"] + self.rms_norm_eps = t["rms_norm_eps"] + self.dense_mlp_idx = int(t.get("dense_mlp_idx", 0)) + self.dense_intermediate_size = int(t.get("dense_intermediate_size", t["intermediate_size"])) + self.logits_mup_width_multiplier = t.get("logits_mup_width_multiplier", None) + self.route_norm = t.get("route_norm", None) + self.use_global_scale = t.get("use_global_scale", False) + # Debug knobs: defaults are the production recipe; MILES_INKLING_* env vars override. + self.attn_backend = os.environ.get("MILES_INKLING_ATTN_BACKEND", "flex") # flex | te | fa4 + self.sconv_impl = os.environ.get("MILES_INKLING_SCONV_IMPL", "triton") # triton | torch + self.sconv_packed = os.environ.get("MILES_INKLING_SCONV_PACKED", "0") == "1" + self.freeze_global_scale = os.environ.get("MILES_INKLING_FREEZE_GLOBAL_SCALE", "all") # all | router | none + + +def build_inkling_config( + text_cfg: dict, + tp=1, + ep=1, + pp=1, + bf16=True, + sp=False, + etp=1, + cp=1, + varlen=True, + permute_fusion=False, + fp32_residual=False, + pp_first_stage_layers=None, + pp_last_stage_layers=None, +) -> TransformerConfig: + inter = text_cfg["intermediate_size"] + ns = text_cfg["n_shared_experts"] + denseI = int(text_cfg.get("dense_intermediate_size", inter)) + cfg = TransformerConfig( + hidden_dropout=0.0, + attention_dropout=0.0, + num_layers=text_cfg["num_hidden_layers"], + hidden_size=text_cfg["hidden_size"], + num_attention_heads=text_cfg["num_attention_heads"], + num_query_groups=text_cfg["num_key_value_heads"], + kv_channels=text_cfg["head_dim"], + ffn_hidden_size=denseI, + tensor_model_parallel_size=tp, + pipeline_model_parallel_size=pp, + num_layers_in_first_pipeline_stage=pp_first_stage_layers, + num_layers_in_last_pipeline_stage=pp_last_stage_layers, + expert_model_parallel_size=ep, + sequence_parallel=sp, + expert_tensor_parallel_size=etp, + context_parallel_size=cp, + variable_seq_lengths=varlen, + fp32_residual_connection=fp32_residual, + moe_permute_fusion=permute_fusion, + num_moe_experts=text_cfg["n_routed_experts"], + moe_router_topk=text_cfg["num_experts_per_tok"], + moe_ffn_hidden_size=inter, + moe_shared_expert_intermediate_size=inter if ns > 0 else None, + moe_router_score_function="sigmoid", + moe_router_enable_expert_bias=True, + moe_router_pre_softmax=True, + moe_router_load_balancing_type="seq_aux_loss", + moe_aux_loss_coeff=0.0, + moe_router_bias_update_rate=0.0, + moe_router_dtype="fp32", + moe_grouped_gemm=True, + moe_token_dispatcher_type="alltoall", + add_bias_linear=False, + normalization="RMSNorm", + layernorm_epsilon=text_cfg["rms_norm_eps"], + qk_layernorm=True, + bf16=bf16, + params_dtype=torch.bfloat16 if bf16 else torch.float32, + gated_linear_unit=True, + pipeline_dtype=torch.bfloat16 if bf16 else torch.float32, + ) + cfg.inkling = InklingExtra(text_cfg) + cfg.moe_activation_in_fp32 = True + cfg.moe_combine_in_fp32 = True + _didx = cfg.inkling.dense_mlp_idx + cfg.moe_layer_freq = [0] * _didx + [1] * (cfg.num_layers - _didx) + cfg.hetereogenous_dist_checkpoint = True + return cfg + + +def get_inkling_layer_spec(config) -> ModuleSpec: + spec = get_gpt_layer_with_transformer_engine_spec( + num_experts=config.num_moe_experts, moe_grouped_gemm=config.moe_grouped_gemm, qk_layernorm=False + ) + attn = spec.submodules.self_attention + spec.submodules.self_attention = ModuleSpec( + module=InklingSelfAttention, params={"attn_mask_type": AttnMaskType.causal}, submodules=attn.submodules + ) + mlp = spec.submodules.mlp + moe_subs = mlp.submodules + moe_subs.router = InklingRouter + moe_subs.shared_experts = ModuleSpec(module=InklingSharedExperts, submodules=moe_subs.shared_experts.submodules) + spec.submodules.mlp = ModuleSpec( + module=InklingMoELayer, params=getattr(mlp, "params", None) or {}, submodules=moe_subs + ) + return spec + + +def get_inkling_dense_layer_spec(config) -> ModuleSpec: + spec = get_gpt_layer_with_transformer_engine_spec(qk_layernorm=False) + attn = spec.submodules.self_attention + spec.submodules.self_attention = ModuleSpec( + module=InklingSelfAttention, params={"attn_mask_type": AttnMaskType.causal}, submodules=attn.submodules + ) + mlp = spec.submodules.mlp + spec.submodules.mlp = ModuleSpec( + module=InklingDenseMLP, params=getattr(mlp, "params", None) or {}, submodules=mlp.submodules + ) + return spec + + +def get_inkling_block_spec(config, vp_stage=None): + from megatron.core.transformer.transformer_block import TransformerBlockSubmodules, _get_block_submodules + from megatron.core.transformer.transformer_layer import get_transformer_layer_offset + + base = _get_block_submodules(config, get_inkling_layer_spec(config), vp_stage) + dense_idx = getattr(config.inkling, "dense_mlp_idx", 0) + if dense_idx <= 0: + return base + offset = get_transformer_layer_offset(config, vp_stage=vp_stage) + local = [ + (get_inkling_dense_layer_spec(config) if (offset + i) < dense_idx else get_inkling_layer_spec(config)) + for i in range(len(base.layer_specs)) + ] + return TransformerBlockSubmodules(layer_specs=local, layer_norm=base.layer_norm) + + +def get_inkling_spec(args, config, vp_stage=None): + """--spec entry for the miles standard provider path.""" + import json + + text_cfg = json.load(open(f"{args.hf_checkpoint}/config.json"))["text_config"] + config.inkling = InklingExtra(text_cfg) + config.moe_router_dtype = "fp32" + config.moe_router_bias_update_rate = 0.0 + config.moe_router_topk_scaling_factor = None + config.moe_activation_in_fp32 = True + config.moe_combine_in_fp32 = True + return get_inkling_layer_spec(config) + + +class InklingGPTModel(GPTModel): + def __init__(self, *args, **kw): + super().__init__(*args, **kw) + if getattr(self, "pre_process", False) and getattr(self.config.inkling, "use_embed_norm", False): + emb = self.embedding + emb.embed_norm = te.RMSNorm( + self.config.hidden_size, eps=self.config.inkling.rms_norm_eps, params_dtype=self.config.params_dtype + ) + _orig_emb_forward = emb.forward + + _fp32res = bool(getattr(self.config, "fp32_residual_connection", False)) + + def _emb_forward(*a, _orig=_orig_emb_forward, _norm=emb.embed_norm, _fp32=_fp32res, **k): + out = _orig(*a, **k) + if _fp32: + h = out.float() + h = h * torch.rsqrt(h.pow(2).mean(-1, keepdim=True) + getattr(_norm, "eps", 1e-6)) + return (h * _norm.weight.float()).bfloat16().float() + return _norm(out) + + emb.forward = _emb_forward + if getattr(self, "post_process", False) and bool(getattr(self.config, "fp32_residual_connection", False)): + fln = getattr(self.decoder, "final_layernorm", None) + if fln is not None: + _orig_fln = fln.forward + _pdt = self.config.params_dtype + + def _fln_forward(x, *a, _orig=_orig_fln, _dt=_pdt, **k): + return _orig(x.to(_dt) if x.dtype != _dt else x, *a, **k) + + fln.forward = _fln_forward + mup = getattr(self.config.inkling, "logits_mup_width_multiplier", None) + if getattr(self, "post_process", False) and mup: + _ol = self.output_layer + _orig_ol_forward = _ol.forward + _mup = float(mup) + + def _ol_forward(x, *a, _orig=_orig_ol_forward, _m=_mup, **k): + return _orig(x / _m, *a, **k) + + _ol.forward = _ol_forward + self._freeze_global_scale() + + def _freeze_global_scale(self): + """Freeze per-layer global_scale params per config.inkling.freeze_global_scale + (all | router | none; router = the MoE gate scale only).""" + mode = getattr(self.config.inkling, "freeze_global_scale", "all") + if mode == "none": + return + n = 0 + for name, p in self.named_parameters(): + if name.rsplit(".", 1)[-1] != "global_scale": + continue + if mode == "all" or (mode == "router" and "router.global_scale" in name): + if p.requires_grad: + p.requires_grad = False + n += 1 + if n: + logger.info( + f"inkling freeze_global_scale={mode}: froze {n} global_scale params " + f"(requires_grad=False; excluded from optimizer + grad norm)" + ) + + +def inkling_model_provider(pre_process=True, post_process=True, vp_stage=None, *, mm_towers=False): + import json + + from megatron.training import get_args + + args = get_args() + if getattr(args, "context_parallel_size", 1) > 1: + assert getattr(args, "allgather_cp", False), "Inkling CP requires --allgather-cp (zigzag CP not supported)" + text_cfg = json.load(open(f"{args.hf_checkpoint}/config.json"))["text_config"] + config = build_inkling_config( + text_cfg, + tp=args.tensor_model_parallel_size, + ep=args.expert_model_parallel_size, + pp=args.pipeline_model_parallel_size, + bf16=args.bf16, + sp=args.sequence_parallel, + etp=getattr(args, "expert_tensor_parallel_size", 1) or 1, + cp=getattr(args, "context_parallel_size", 1) or 1, + varlen=getattr(args, "variable_seq_lengths", True), + permute_fusion=getattr(args, "moe_permute_fusion", False), + fp32_residual=getattr(args, "fp32_residual_connection", False), + pp_first_stage_layers=getattr(args, "decoder_first_pipeline_num_layers", None), + pp_last_stage_layers=getattr(args, "decoder_last_pipeline_num_layers", None), + ) + model = InklingGPTModel( + config=config, + transformer_layer_spec=get_inkling_block_spec(config, vp_stage=vp_stage), + vocab_size=text_cfg["vocab_size"], + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + position_embedding_type="none", + share_embeddings_and_output_weights=False, + vp_stage=vp_stage, + ) + if mm_towers: + from miles_plugins.models.inkling.mm_towers import wire_mm_towers + + wire_mm_towers(model, args.hf_checkpoint) + return model + + +def inkling_mm_model_provider(pre_process=True, post_process=True, vp_stage=None): + """Multimodal provider: the text model plus the frozen HF vision/audio towers. + + A separate entry point instead of a CLI switch -- multimodal launch scripts pass + this as --custom-model-provider-path. + """ + return inkling_model_provider(pre_process, post_process, vp_stage, mm_towers=True) diff --git a/miles_plugins/models/inkling/ops.py b/miles_plugins/models/inkling/ops.py new file mode 100644 index 00000000000..ef056d4d34a --- /dev/null +++ b/miles_plugins/models/inkling/ops.py @@ -0,0 +1,625 @@ +"""Inkling custom ops in one place: CP/SP helpers, precision-aligned fp32 ops with their +triton kernels (forward bit-identical to the sglang serving kernels), and the attention +backends (flex / TE reference / FA4 with Inkling's relative position bias).""" + +from collections.abc import Callable +from functools import cache + +import megatron.core.parallel_state as ps +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from megatron.core.tensor_parallel.mappings import ( + gather_from_sequence_parallel_region, + scatter_to_sequence_parallel_region, +) + + +try: + from torch.nn.attention.flex_attention import create_block_mask, flex_attention +except Exception: # pragma: no cover + flex_attention = create_block_mask = None + + +def cp_world(): + """(cp_size, cp_rank, cp_group) — cp_size 1 means no context parallelism.""" + cp = ps.get_context_parallel_world_size() + if cp <= 1: + return 1, 0, None + return cp, ps.get_context_parallel_rank(), ps.get_context_parallel_group() + + +def cp_all_gather(x, group, world): + """all-gather [s, ...] -> [s*world, ...] in CP-rank order. Assumes CONTIGUOUS CP sharding + (rank r holds tokens [r*s,(r+1)*s)) -> REQUIRES --allgather-cp (miles' default CP layout is + zigzag/load-balanced; the provider asserts allgather_cp when cp>1).""" + xs = [torch.empty_like(x) for _ in range(world)] + torch.distributed.all_gather(xs, x.contiguous(), group=group) + return torch.cat(xs, dim=0) + + +def inkling_sp_residual_conv(config, conv, x_sbh, seqlens): + """Residual depthwise sconv on [s,b,h]. Under SP/CP the sequence is sharded, so a local-shard + causal conv misses the left context -> gather -> conv on the full sequence -> take this rank's + slice. Exact because the conv is residual. seqlens must be the FULL-sequence segment lengths.""" + sp = getattr(config, "sequence_parallel", False) and ps.get_tensor_model_parallel_world_size() > 1 + cp, cp_rank, cp_group = cp_world() + x = gather_from_sequence_parallel_region(x_sbh, tensor_parallel_output_grad=False) if sp else x_sbh + if cp > 1: + x = cp_all_gather(x, cp_group, cp) + s, b, h = x.shape + x = conv(x.reshape(s * b, h), seqlens).reshape(s, b, h) + if cp > 1: + sloc = s // cp + x = x[cp_rank * sloc : (cp_rank + 1) * sloc] + return scatter_to_sequence_parallel_region(x) if sp else x + + +def seqlens_from_packed(packed_seq_params, T): + """THD packing: per-segment token lengths (over the full post-SP-gather length T) from + cu_seqlens_q, so attention/sconv/rel-bias never cross packed-sequence boundaries. Clip to T: + keep whole segments, split the boundary one, trailing pad as its own segment.""" + if packed_seq_params is None or getattr(packed_seq_params, "cu_seqlens_q", None) is None: + return None + cu = packed_seq_params.cu_seqlens_q + raw = [int(s) for s in (cu[1:] - cu[:-1]).tolist() if s > 0] + seqlens, acc = [], 0 + for s in raw: + if acc + s <= T: + seqlens.append(s) + acc += s + else: + if T - acc > 0: + seqlens.append(T - acc) + acc = T + break + if acc < T: + seqlens.append(T - acc) + return seqlens + + +@triton.jit +def _swiglu_fwd_kernel( + x_ptr, + scale_ptr, + y_ptr, + M, + N, + HAS_SCALE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + mask_m = offs_m < M + mask_n = offs_n < N + mask = mask_m[:, None] & mask_n[None, :] + g = tl.load(x_ptr + offs_m[:, None] * (2 * N) + offs_n[None, :], mask=mask, other=0.0).to(tl.float32) + u = tl.load(x_ptr + offs_m[:, None] * (2 * N) + N + offs_n[None, :], mask=mask, other=0.0).to(tl.float32) + y = g * tl.sigmoid(g) * u + if HAS_SCALE: + sc = tl.load(scale_ptr + offs_m, mask=mask_m, other=0.0).to(tl.float32) + y = y * sc[:, None] + tl.store(y_ptr + offs_m[:, None] * N + offs_n[None, :], y, mask=mask) + + +@triton.jit +def _swiglu_bwd_kernel( + x_ptr, + scale_ptr, + go_ptr, + dx_ptr, + dscale_ptr, + M, + N, + HAS_SCALE: tl.constexpr, + NEED_DSCALE: tl.constexpr, + BLOCK_N: tl.constexpr, +): + row = tl.program_id(0) + if row >= M: + return + sc = 1.0 + if HAS_SCALE: + sc = tl.load(scale_ptr + row).to(tl.float32) + acc = tl.zeros((BLOCK_N,), dtype=tl.float32) + for n0 in tl.range(0, N, BLOCK_N): + offs_n = n0 + tl.arange(0, BLOCK_N) + mask_n = offs_n < N + g = tl.load(x_ptr + row * (2 * N) + offs_n, mask=mask_n, other=0.0).to(tl.float32) + u = tl.load(x_ptr + row * (2 * N) + N + offs_n, mask=mask_n, other=0.0).to(tl.float32) + go = tl.load(go_ptr + row * N + offs_n, mask=mask_n, other=0.0).to(tl.float32) + s = tl.sigmoid(g) + silu = g * s + gos = go * sc + d_g = gos * u * (s + silu * (1 - s)) + d_u = gos * silu + tl.store(dx_ptr + row * (2 * N) + offs_n, d_g, mask=mask_n) + tl.store(dx_ptr + row * (2 * N) + N + offs_n, d_u, mask=mask_n) + if NEED_DSCALE: + acc += tl.where(mask_n, go * silu * u, 0.0) + if NEED_DSCALE: + tl.store(dscale_ptr + row, tl.sum(acc)) + + +class _SwigluFP32Triton(torch.autograd.Function): + @staticmethod + def forward(ctx, fc1_out: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor: + shp = fc1_out.shape + x = fc1_out.reshape(-1, shp[-1]).contiguous() + M, N2 = x.shape + N = N2 // 2 + sc = scale.reshape(-1).contiguous().float() if scale is not None else x.new_empty(1, dtype=torch.float32) + y = torch.empty(M, N, device=x.device, dtype=x.dtype) + if M > 0: + grid = (triton.cdiv(M, 32), triton.cdiv(N, 128)) + _swiglu_fwd_kernel[grid](x, sc, y, M, N, HAS_SCALE=scale is not None, BLOCK_M=32, BLOCK_N=128) + ctx.save_for_backward(fc1_out, scale) + return y.reshape(*shp[:-1], N) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + fc1_out, scale = ctx.saved_tensors + shp = fc1_out.shape + x = fc1_out.reshape(-1, shp[-1]).contiguous() + M, N2 = x.shape + N = N2 // 2 + go = grad_out.reshape(-1, N).contiguous() + sc = scale.reshape(-1).contiguous().float() if scale is not None else x.new_empty(1, dtype=torch.float32) + need_dscale = scale is not None and ctx.needs_input_grad[1] + dx = torch.empty_like(x) + dsc = ( + torch.empty(M, device=x.device, dtype=torch.float32) + if need_dscale + else x.new_empty(1, dtype=torch.float32) + ) + if M > 0: + _swiglu_bwd_kernel[(M,)]( + x, + sc, + go, + dx, + dsc, + M, + N, + HAS_SCALE=scale is not None, + NEED_DSCALE=need_dscale, + BLOCK_N=256, + ) + d_scale = dsc.reshape(scale.shape).to(scale.dtype) if need_dscale else None + return dx.reshape(shp), d_scale + + +def swiglu_fp32_triton(fc1_out: torch.Tensor, per_token_scale: torch.Tensor | None = None) -> torch.Tensor: + """Triton swiglu, forward bit-identical to sglang's silu_and_mul kernel.""" + return _SwigluFP32Triton.apply(fc1_out, per_token_scale) + + +@triton.jit +def _sconv_fwd_kernel( + x_ptr, + w_ptr, + bos_ptr, + y_ptr, + T, + D, + W: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_D: tl.constexpr, +): + t_off = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + d_off = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + t_mask = t_off < T + d_mask = d_off < D + td_mask = t_mask[:, None] & d_mask[None, :] + + bos = tl.load(bos_ptr + t_off, mask=t_mask, other=0).to(tl.int64) + acc = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32) + x_cur = tl.load(x_ptr + t_off[:, None] * D + d_off[None, :], mask=td_mask, other=0) + + for iw in tl.static_range(W): + shifted = t_off - (W - 1) + iw + if iw == W - 1: + tap = x_cur.to(tl.float32) + else: + in_x = shifted >= bos + x_val = tl.load( + x_ptr + shifted[:, None] * D + d_off[None, :], + mask=in_x[:, None] & td_mask, + other=0, + eviction_policy="evict_last", + ) + tap = x_val.to(tl.float32) + w_val = tl.load(w_ptr + d_off * W + iw, mask=d_mask, other=0).to(tl.float32) + acc += tap * w_val[None, :] + + acc += x_cur.to(tl.float32) + tl.store(y_ptr + t_off[:, None] * D + d_off[None, :], acc.to(y_ptr.dtype.element_ty), mask=td_mask) + + +@triton.jit +def _sconv_bwd_dx_kernel( + go_ptr, + w_ptr, + eos_ptr, + dx_ptr, + T, + D, + W: tl.constexpr, + BLOCK_T: tl.constexpr, + BLOCK_D: tl.constexpr, +): + t_off = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T) + d_off = tl.program_id(1) * BLOCK_D + tl.arange(0, BLOCK_D) + t_mask = t_off < T + d_mask = d_off < D + td_mask = t_mask[:, None] & d_mask[None, :] + + eos = tl.load(eos_ptr + t_off, mask=t_mask, other=0).to(tl.int64) + acc = tl.zeros([BLOCK_T, BLOCK_D], dtype=tl.float32) + + for iw in tl.static_range(W): + fwd_t = t_off + (W - 1) - iw + in_seg = fwd_t < eos + g_val = tl.load( + go_ptr + fwd_t[:, None] * D + d_off[None, :], + mask=in_seg[:, None] & td_mask, + other=0, + eviction_policy="evict_last", + ).to(tl.float32) + w_val = tl.load(w_ptr + d_off * W + iw, mask=d_mask, other=0).to(tl.float32) + acc += g_val * w_val[None, :] + + g_cur = tl.load(go_ptr + t_off[:, None] * D + d_off[None, :], mask=td_mask, other=0).to(tl.float32) + acc += g_cur + tl.store(dx_ptr + t_off[:, None] * D + d_off[None, :], acc.to(dx_ptr.dtype.element_ty), mask=td_mask) + + +def _bos_eos(seqlens, T, device): + if seqlens is None or len(seqlens) <= 1: + bos = torch.zeros(T, device=device, dtype=torch.int32) + eos = torch.full((T,), T, device=device, dtype=torch.int32) + return bos, eos + sl = torch.as_tensor(list(seqlens), device=device, dtype=torch.int32) + ends = sl.cumsum(0) + starts = ends - sl + bos = torch.repeat_interleave(starts, sl) + eos = torch.repeat_interleave(ends, sl) + return bos.to(torch.int32), eos.to(torch.int32) + + +class _SconvFP32Triton(torch.autograd.Function): + @staticmethod + def forward(ctx, x: torch.Tensor, weight: torch.Tensor, seqlens) -> torch.Tensor: + T, D = x.shape + W = weight.shape[-1] + xc = x.contiguous() + w2 = weight.reshape(D, W).contiguous() + bos, eos = _bos_eos(seqlens, T, x.device) + y = torch.empty_like(xc) + if T > 0: + grid = (triton.cdiv(T, 64), triton.cdiv(D, 128)) + _sconv_fwd_kernel[grid](xc, w2, bos, y, T, D, W=W, BLOCK_T=64, BLOCK_D=128) + ctx.save_for_backward(xc, w2, bos, eos) + ctx.w_shape = weight.shape + return y + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + xc, w2, bos, eos = ctx.saved_tensors + T, D = xc.shape + W = w2.shape[-1] + go = grad_out.contiguous() + dx = torch.empty_like(xc) + if T > 0: + grid = (triton.cdiv(T, 64), triton.cdiv(D, 128)) + _sconv_bwd_dx_kernel[grid](go, w2, eos, dx, T, D, W=W, BLOCK_T=64, BLOCK_D=128) + go32 = go.float() + x32 = xc.float() + dw = torch.empty(D, W, device=xc.device, dtype=torch.float32) + ar = torch.arange(T, device=xc.device, dtype=torch.int32) + for iw in range(W): + shift = (W - 1) - iw + if shift == 0: + dw[:, iw] = (go32 * x32).sum(0) + else: + valid = (ar[shift:] - shift) >= bos[shift:] + dw[:, iw] = (go32[shift:] * x32[: T - shift] * valid.unsqueeze(1)).sum(0) + return dx, dw.reshape(ctx.w_shape).to(w2.dtype), None + + +def sconv_fp32_triton(x: torch.Tensor, weight: torch.Tensor, seqlens=None) -> torch.Tensor: + """Packed-sequence residual sconv, forward bit-identical to sglang's kernel.""" + return _SconvFP32Triton.apply(x, weight, seqlens) + + +def _maybe_compile(fn): + return torch.compile(fn, dynamic=True) + + +@_maybe_compile +def _swiglu_fwd(fc1_out, scale): + g, u = torch.chunk(fc1_out.float(), 2, dim=-1) + y = F.silu(g) * u + if scale is not None: + y = y * scale.float() + return y.to(fc1_out.dtype) + + +@_maybe_compile +def _swiglu_bwd(fc1_out, scale, grad_out, need_dscale: bool): + g, u = torch.chunk(fc1_out.float(), 2, dim=-1) + s = torch.sigmoid(g) + silu = g * s + go = grad_out.float() + if scale is not None: + go = go * scale.float() + d_g = go * u * (s + silu * (1 - s)) + d_u = go * silu + d_scale = None + if need_dscale: + d_scale = (grad_out.float() * silu * u).sum(-1, keepdim=True).to(scale.dtype) + return torch.cat([d_g, d_u], dim=-1).to(fc1_out.dtype), d_scale + + +class _SwigluFP32Func(torch.autograd.Function): + @staticmethod + def forward(ctx, fc1_out: torch.Tensor, scale: torch.Tensor | None) -> torch.Tensor: + ctx.save_for_backward(fc1_out, scale) + return _swiglu_fwd(fc1_out, scale) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + fc1_out, scale = ctx.saved_tensors + need_dscale = scale is not None and ctx.needs_input_grad[1] + d_fc1, d_scale = _swiglu_bwd(fc1_out, scale, grad_out, need_dscale) + return d_fc1, d_scale + + +def inkling_swiglu_fp32(fc1_out: torch.Tensor, per_token_scale: torch.Tensor | None = None) -> torch.Tensor: + """fp32 swiglu, one round back to bf16 (triton, bit-identical to serving silu_and_mul).""" + return swiglu_fp32_triton(fc1_out, per_token_scale) + + +@_maybe_compile +def _sum2_fwd(a, b): + return (a.float() + b.float()).to(a.dtype) + + +@_maybe_compile +def _sum3_fwd(a, b, c): + return (a.float() + b.float() + c.float()).to(a.dtype) + + +class _SumFP32Func(torch.autograd.Function): + @staticmethod + def forward(ctx, *ys: torch.Tensor) -> torch.Tensor: + ctx.n = len(ys) + if len(ys) == 2: + return _sum2_fwd(*ys) + if len(ys) == 3: + return _sum3_fwd(*ys) + out = ys[0].float() + for y in ys[1:]: + out = out + y.float() + return out.to(ys[0].dtype) + + @staticmethod + def backward(ctx, grad_out: torch.Tensor): + return (grad_out,) * ctx.n + + +def inkling_sum_fp32(ys: list[torch.Tensor]) -> torch.Tensor: + """Σ ys in fp32, single round back (SGLang `_sum_dim0` parity).""" + return _SumFP32Func.apply(*ys) + + +@_maybe_compile +def _sconv_fwd(x, weight): + C = x.shape[1] + xp = F.pad(x.float().t().unsqueeze(0), (weight.shape[-1] - 1, 0)) + y = x.float() + F.conv1d(xp, weight.float(), groups=C).squeeze(0).t() + return y.to(x.dtype) + + +def inkling_sconv_fp32(x: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Residual depthwise causal conv on [T, C] in fp32, single round back (SGLang sconv parity).""" + return _sconv_fwd(x, weight) + + +def inkling_sconv_fp32_packed( + x: torch.Tensor, weight: torch.Tensor, seqlens, impl: str = "triton", packed: bool = False +) -> torch.Tensor: + """Packed-sequence sconv. impl=triton is bit-identical to serving; packed runs the whole + batch as one full-length conv + exact boundary re-compute instead of per-segment convs.""" + if x.is_cuda and impl == "triton": + return sconv_fp32_triton(x, weight, seqlens) + if seqlens is None or len(seqlens) <= 1: + return inkling_sconv_fp32(x, weight) + if not packed: + return torch.cat([inkling_sconv_fp32(s, weight) for s in x.split(list(seqlens))], 0) + + k = weight.shape[-1] + T, C = x.shape + xf = x.float() + w = weight.float() + xp = F.pad(xf.t().unsqueeze(0), (k - 1, 0)) + y = xf + F.conv1d(xp, w, groups=C).squeeze(0).t() + + if k > 1: + dev = x.device + sl = torch.as_tensor(list(seqlens), device=dev, dtype=torch.long) + starts = sl.cumsum(0)[:-1] + lens = sl[1:] + ds = torch.arange(k - 1, device=dev) + t_idx = starts.view(-1, 1) + ds.view(1, -1) + valid = ds.view(1, -1) < lens.view(-1, 1) + xw = xf[t_idx.clamp(max=T - 1)] + wt = w.squeeze(1).t() + W2 = x.new_zeros(k - 1, k - 1, C, dtype=torch.float32) + for d in range(k - 1): + for i in range(d + 1): + W2[d, i] = wt[k - 1 - (d - i)] + corr = torch.einsum("sic,dic->sdc", xw, W2) + y = y.index_put((t_idx[valid],), (xw + corr)[valid]) + + return y.to(x.dtype) + + +try: + import cutlass.cute as cute + from cutlass.cute import Float32 + from flash_attn.cute.seqlen_info import SeqlenInfoQK +except Exception as _import_error: + cute = None + Float32 = None + SeqlenInfoQK = None + _cute_import_error = _import_error +else: + _cute_import_error = None + + +@cache +def get_inkling_relative_attention_score_mod(rel_extent: int) -> Callable: + if cute is None or Float32 is None or SeqlenInfoQK is None: + raise ImportError( + "Inkling relative attention requires the vendored FA4 CUTE interface." + ) from _cute_import_error + + @cute.jit + def score_mod_rel_bias( + scores: cute.TensorSSA, + b_idx: cute.TensorSSA, + h_idx: cute.TensorSSA, + q_idx: cute.TensorSSA, + kv_idx: cute.TensorSSA, + seqlen_info: SeqlenInfoQK, + aux_tensors: list[cute.Tensor], + ) -> cute.TensorSSA: + rel_logits = aux_tensors[0] + + seqlen_local_offset = seqlen_info.seqlen_k - seqlen_info.seqlen_q + rel_dist = (q_idx + seqlen_local_offset) - kv_idx + global_q_idx = seqlen_info.offset_q + q_idx + + rel_dist_0 = rel_dist[0] + rel_idx = rel_dist_0 if rel_dist_0 >= 0 else 0 + rel_idx = rel_idx if rel_idx < rel_extent else (rel_extent - 1) + + rel_bias = rel_logits[global_q_idx[0], h_idx[0], rel_idx] + rel_bias = Float32(rel_bias) if rel_dist_0 == rel_idx else Float32(0.0) + return scores + rel_bias + + return score_mod_rel_bias + + +_FLEX_COMPILED = None + + +def flex_compiled(): + global _FLEX_COMPILED + if _FLEX_COMPILED is None: + assert flex_attention is not None, "flex_attention needs torch>=2.5 + inductor" + torch._dynamo.config.cache_size_limit = 1024 + torch._dynamo.config.accumulated_cache_size_limit = 1024 + _FLEX_COMPILED = torch.compile(flex_attention, dynamic=True, mode="max-autotune-no-cudagraphs") + return _FLEX_COMPILED + + +def inkling_te_attention(q, k, v, rel_logits, seqlens, window_left, is_local, scale): + import transformer_engine.pytorch as te + + T, nh, hd = q.shape + nkv = k.shape[1] + RE = rel_logits.shape[-1] + dpa = te.DotProductAttention( + num_attention_heads=nh, + kv_channels=hd, + num_gqa_groups=nkv, + attention_dropout=0.0, + qkv_format="sbhd", + softmax_scale=scale, + ) + rb = rel_logits.permute(1, 0, 2) + qi = torch.arange(T, device=q.device).view(T, 1) + ki = torch.arange(T, device=q.device).view(1, T) + rd = qi - ki + valid = (rd >= 0) & (rd < RE) + idx = rd.clamp(0, RE - 1) + bias = (torch.gather(rb, 2, idx.unsqueeze(0).expand(nh, T, T)) * valid.unsqueeze(0)).unsqueeze(0) + win = (window_left, 0) if is_local else (-1, -1) + if seqlens is not None and len(seqlens) > 1: + seg = torch.repeat_interleave( + torch.arange(len(seqlens), device=q.device), torch.tensor(seqlens, device=q.device) + ) + cross = (seg.view(T, 1) != seg.view(1, T)).view(1, 1, T, T) + bias = bias.masked_fill(cross, -1e9) + ctx = dpa( + q.unsqueeze(1).contiguous(), + k.unsqueeze(1).contiguous(), + v.unsqueeze(1).contiguous(), + attn_mask_type="causal", + window_size=win, + core_attention_bias_type="post_scale_bias", + core_attention_bias=bias.to(q.dtype), + ) + return ctx.reshape(T, nh, hd) + + +def _fa4_fwd(q, k, v, rel_logits, seqlens, window_left, is_local, scale, rel_extent): + from flash_attn.cute.interface import flash_attn_varlen_func + + T = q.shape[0] + sl = list(seqlens) if seqlens else [T] + cu = torch.zeros(len(sl) + 1, device=q.device, dtype=torch.int32) + cu[1:] = torch.tensor(sl, device=q.device, dtype=torch.int32).cumsum(0) + win = (window_left, 0) if is_local else (-1, -1) + out = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=cu, + cu_seqlens_k=cu, + max_seqlen_q=max(sl), + max_seqlen_k=max(sl), + softmax_scale=scale, + causal=True, + window_size=win, + score_mod=get_inkling_relative_attention_score_mod(rel_extent), + aux_tensors=[rel_logits.contiguous().float()], + ) + if isinstance(out, tuple): + out = out[0] + return out + + +class _InklingFA4Attention(torch.autograd.Function): + @staticmethod + def forward(ctx, q, k, v, rel_logits, seqlens, window_left, is_local, scale, rel_extent): + with torch.no_grad(): + out = _fa4_fwd(q, k, v, rel_logits, seqlens, window_left, is_local, scale, rel_extent) + ctx.save_for_backward(q, k, v, rel_logits) + ctx.meta = (tuple(seqlens) if seqlens else None, window_left, is_local, scale, rel_extent) + return out + + @staticmethod + def backward(ctx, grad_out): + q, k, v, rel_logits = ctx.saved_tensors + seqlens, window_left, is_local, scale, rel_extent = ctx.meta + sl = list(seqlens) if seqlens else None + with torch.enable_grad(): + qd = q.detach().requires_grad_(True) + kd = k.detach().requires_grad_(True) + vd = v.detach().requires_grad_(True) + rd = rel_logits.detach().requires_grad_(True) + y = inkling_te_attention(qd, kd, vd, rd, sl, window_left, is_local, scale) + gq, gk, gv, gr = torch.autograd.grad(y, [qd, kd, vd, rd], grad_out) + return gq, gk, gv, gr, None, None, None, None, None + + +def inkling_fa4_attention(q, k, v, rel_logits, seqlens, window_left, is_local, scale, rel_extent): + """q [T,nh,hd] / k,v [T,nkv,hd] / rel_logits [T,nh,RE] fp32 -> [T,nh,hd]; fwd = sglang FA4, bwd recompute.""" + return _InklingFA4Attention.apply(q, k, v, rel_logits, seqlens, window_left, is_local, scale, rel_extent) diff --git a/miles_plugins/models/inkling/vision_encoder.py b/miles_plugins/models/inkling/vision_encoder.py new file mode 100644 index 00000000000..6db637a9fa4 --- /dev/null +++ b/miles_plugins/models/inkling/vision_encoder.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +from typing import cast + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + + +class InklingVisionConfig: + """Duck-typed stand-in for the serving-side config (fields the encoder reads).""" + + def __init__(self, d: dict): + self.decoder_dmodel = d["decoder_dmodel"] + self.n_channels = d.get("n_channels", 3) + self.n_layers = d["n_layers"] + self.patch_size = d["patch_size"] + self.temporal_patch_size = d.get("temporal_patch_size", 1) + self.use_vision_norm = d.get("use_vision_norm", True) + + +class RMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.hidden_size = hidden_size + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.numel() == 0: + return x + return F.rms_norm(x.float(), (self.hidden_size,), self.weight.float(), self.variance_epsilon).to(x.dtype) + + +def _prime_factors(n: int) -> list[int]: + """ + Return the prime factors of *n* in ascending order. + """ + if n < 1: + raise ValueError("n must be a positive integer") + + factors: list[int] = [] + + while n % 2 == 0: + factors.append(2) + n //= 2 + + p = 3 + while p * p <= n: + while n % p == 0: + factors.append(p) + n //= p + p += 2 + + if n > 1: + factors.append(n) + return factors + + +def plan_out_scales( + temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3 +) -> list[tuple[int, int, int, int]]: + if patch_size <= 1: + raise ValueError("patch_size must be greater than 1, otherwise this doesn't make sense") + + def _round_up(x: int) -> int: + return int(np.ceil(x / 64)) * 64 + + last_h_scale = 1 + scales: list[tuple[int, int, int, int]] = [(1, 1, 1, n_channels)] + for pscale in _prime_factors(patch_size)[::-1]: + last_h_scale *= pscale + scales.append( + ( + 1, + last_h_scale, + last_h_scale, + _round_up((last_h_scale**2) * n_channels), + ) + ) + last_t_scale = 1 + for tscale in _prime_factors(temporal_patch_size)[::-1]: + last_t_scale *= tscale + scales.append( + ( + last_t_scale, + last_h_scale, + last_h_scale, + _round_up((last_h_scale**2) * n_channels * last_t_scale), + ) + ) + + size_reduction = np.prod(np.array(scales)[:, :-1], 1) + + log_ideal_scales = np.linspace( + 0, + np.log(patch_size * patch_size * temporal_patch_size * n_channels), + n_layers + 1, + ) + cost_matrix = np.abs(log_ideal_scales[:, None] - np.log(size_reduction)[None]) + + if n_layers >= len(scales): + idxs = np.argmin(cost_matrix, axis=1) + else: + from scipy.optimize import linear_sum_assignment + + idxs = linear_sum_assignment(cost_matrix)[1] + + assert len(idxs) >= 2 + idxs[0] = 0 + idxs[-1] = len(scales) - 1 + + return [scales[i] for i in idxs] + + +def fold_timespace_to_depth(vision_patches_bthwc: torch.Tensor, t_fold: int, hw_fold: int) -> torch.Tensor: + """ + Convert a tensor of shape (B, T, H, W, C) to a tensor of shape (B, T // t, H // hw, W // hw, C * (t * hw**2)) + """ + B, T, H, W, C = vision_patches_bthwc.shape + + assert T % t_fold == 0, f"Temporal dimension {T} must be divisible by {t_fold}" + assert H % hw_fold == 0, f"Height dimension {H} must be divisible by {hw_fold}" + assert W % hw_fold == 0, f"Width dimension {W} must be divisible by {hw_fold}" + + t_new = T // t_fold + h_new = H // hw_fold + w_new = W // hw_fold + + x = vision_patches_bthwc.reshape(B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C) + + x = x.permute(0, 1, 3, 5, 2, 4, 6, 7) + + x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C) + + return x + + +class HMLPPatchEncoder(nn.Module): + def __init__( + self, + config: InklingVisionConfig, + ): + super().__init__() + self.decoder_dmodel = config.decoder_dmodel + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.n_channels = config.n_channels + self.n_layers = config.n_layers + self.use_vision_norm = config.use_vision_norm + + self.scales: list[tuple[int, int, int, int]] = plan_out_scales( + self.temporal_patch_size, self.patch_size, self.n_layers, self.n_channels + ) + self.layers: nn.ModuleDict = nn.ModuleDict() + for i, (start_scale, end_scale) in enumerate(zip(self.scales[:-1], self.scales[1:], strict=False)): + shuffle_mult = ( + (end_scale[0] // start_scale[0]) * (end_scale[1] // start_scale[1]) * (end_scale[2] // start_scale[2]) + ) + if i == self.n_layers - 1: + self.layers[f"linear_{i}"] = nn.Linear(start_scale[3] * shuffle_mult, self.decoder_dmodel, bias=False) + else: + self.layers[f"linear_{i}"] = nn.Linear(start_scale[3] * shuffle_mult, end_scale[3], bias=False) + self.layers[f"norm_{i}"] = RMSNorm(end_scale[3]) + + self.final_norm: RMSNorm | None = None + if self.use_vision_norm: + self.final_norm = RMSNorm(self.decoder_dmodel) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + num_patches, T, H, W, C = x.shape + for i, (start_scale, end_scale) in enumerate(zip(self.scales[:-1], self.scales[1:], strict=False)): + t_fold = end_scale[0] // start_scale[0] + hw_fold = end_scale[1] // start_scale[1] + if hw_fold > 1 or t_fold > 1: + x = fold_timespace_to_depth(x, t_fold, hw_fold) + assert x.shape[1:-1] == ( + T // end_scale[0], + H // end_scale[1], + W // end_scale[2], + ) + x = self.layers[f"linear_{i}"](x) + if i < self.n_layers - 1: + norm = cast(RMSNorm, self.layers[f"norm_{i}"]) + x = norm(x) + x = F.gelu(x) + + if self.final_norm is not None: + x = self.final_norm(x) + + x = x.reshape(num_patches, -1) + return x diff --git a/pyproject.toml b/pyproject.toml index ad1631f277b..fe01b30790b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,6 +35,7 @@ ignore = [ "E501", # Line too long # TODO handle it later ] + [tool.pytest.ini_options] # durations=0 will display all tests execution time, sorted in ascending order starting from from the slowest one. # -vv will also display tests with duration = 0.00s diff --git a/scripts/models/inkling-small.sh b/scripts/models/inkling-small.sh new file mode 100644 index 00000000000..10131973de1 --- /dev/null +++ b/scripts/models/inkling-small.sh @@ -0,0 +1,43 @@ +NLAYERS="${MODEL_ARGS_NUM_LAYERS:-42}" + +# Inkling-Small 276B config (42 layers; derived from the HF config the same way +# inkling.sh maps the Inkling one). +MODEL_ARGS=( + --disable-bias-linear + --num-layers $NLAYERS + --hidden-size 4096 + --ffn-hidden-size 2048 + --num-attention-heads 32 + --group-query-attention + --num-query-groups 8 + --kv-channels 128 + --normalization RMSNorm + --norm-epsilon 1e-6 + --swiglu + --untie-embeddings-and-output-weights + --vocab-size 201024 + --hidden-dropout 0.0 + --attention-dropout 0.0 + --attention-softmax-in-fp32 + --position-embedding-type none + --no-rope-fusion + --no-masked-softmax-fusion + --max-position-embeddings 1048576 + + # MoE + --num-experts 256 + --moe-ffn-hidden-size 2048 + --moe-router-topk 6 + --moe-shared-expert-intermediate-size 2048 + --moe-router-pre-softmax + --moe-router-score-function sigmoid + --moe-router-enable-expert-bias + --moe-router-load-balancing-type seq_aux_loss + --moe-token-dispatcher-type alltoall + --moe-aux-loss-coeff 0 + --moe-grouped-gemm + --qk-layernorm + + # Inkling model provider + --custom-model-provider-path miles_plugins.models.inkling.model.inkling_model_provider +) diff --git a/scripts/models/inkling.sh b/scripts/models/inkling.sh new file mode 100644 index 00000000000..313f4a3db88 --- /dev/null +++ b/scripts/models/inkling.sh @@ -0,0 +1,42 @@ +NLAYERS="${MODEL_ARGS_NUM_LAYERS:-66}" + +# Inkling config (66L full model; set MODEL_ARGS_NUM_LAYERS=4 for the 4-layer slice) +MODEL_ARGS=( + --disable-bias-linear + --num-layers $NLAYERS + --hidden-size 6144 + --ffn-hidden-size 3072 + --num-attention-heads 64 + --group-query-attention + --num-query-groups 8 + --kv-channels 128 + --normalization RMSNorm + --norm-epsilon 1e-6 + --swiglu + --untie-embeddings-and-output-weights + --vocab-size 201024 + --hidden-dropout 0.0 + --attention-dropout 0.0 + --attention-softmax-in-fp32 + --position-embedding-type none + --no-rope-fusion + --no-masked-softmax-fusion + --max-position-embeddings 1048576 + + # MoE + --num-experts 256 + --moe-ffn-hidden-size 3072 + --moe-router-topk 6 + --moe-shared-expert-intermediate-size 3072 + --moe-router-pre-softmax + --moe-router-score-function sigmoid + --moe-router-enable-expert-bias + --moe-router-load-balancing-type seq_aux_loss + --moe-token-dispatcher-type alltoall + --moe-aux-loss-coeff 0 + --moe-grouped-gemm + --qk-layernorm + + # Inkling model provider + --custom-model-provider-path miles_plugins.models.inkling.model.inkling_model_provider +) diff --git a/scripts/run_inkling.py b/scripts/run_inkling.py new file mode 100644 index 00000000000..c7376dae424 --- /dev/null +++ b/scripts/run_inkling.py @@ -0,0 +1,413 @@ +""" +Inkling family training script (Inkling / Inkling-Small / 4-layer slice). + +Supports: + - Inkling 66-layer MoE (frozen vision/audio towers optional). + Verified profiles: 16 nodes x 4 GPUs (TP4 PP4 EP16) and + 12 nodes x 4 GPUs (TP4 PP3 EP16) on GB300. + - Inkling-4layer 4-layer slice for single-node smoke testing. + - Inkling-Small-4layer 4-layer slice of Inkling-Small (fits the 4-GPU CI lane). + - Inkling-Small 42-layer 276B MoE. Verified profile: 4 nodes x 8 GPUs + (TP4 SP PP8 EP4, ctx 4096 / response 2048) on H200: + --lr 5e-5, --rollout-batch-size 64 --global-batch-size 128 + (reward +0.006/rollout). + +Full-parameter GRPO: the paused training actor (params, grads, optimizer +state) spills to node-local NVMe via torch_memory_saver +(--offload-train-target disk). + +Topologies: + - colocate (default) engines share the training GPUs; the actor sleeps + during generation. + - --fully-async disaggregated: --rollout-num-nodes host one 16-GPU + engine while the rest train (train_async.py driver, + generation runs continuously in a background worker). + Verified profile: Inkling on 12 nodes x 4 GPUs as + 8 train + 4 rollout (TP4 PP2 EP16), bf16 grad reduce, + optimizer state streamed through NVMe. Does not + serve eval. + +Tasks: + - dapo_math text (dapo-math-17k), chat template, aime25 eval. + - geo3k vision (geo3k), structured-message rendering, frozen towers, + geo3k-val eval. + +Usage patterns: + + 1. Train on pre-staged checkpoints: + python scripts/run_inkling.py train \ + --model-name Inkling --task dapo_math \ + --num-nodes 16 --num-gpus-per-node 4 + + 2. Individual steps (rsync shared -> node-local NVMe, then train): + python scripts/run_inkling.py prepare-cp --model-name Inkling + python scripts/run_inkling.py train --model-name Inkling ... + + 3. One-shot (prepare-cp when model_local_dir differs, then train): + python scripts/run_inkling.py full-train --model-name Inkling ... + + 4. Fully-async disaggregated training: + MILES_SCRIPT_NUM_NODES=12 python scripts/run_inkling.py train \ + --model-name Inkling --fully-async --rollout-num-nodes 4 \ + --num-gpus-per-node 4 --num-rollout 100 +""" + +import os +from dataclasses import dataclass, field +from typing import Literal + +import typer + +import miles.utils.external_utils.command_utils as U + +app = typer.Typer() + +# model name -> scripts/models/.sh; the 4-layer slices reuse the base +# definition with MODEL_ARGS_NUM_LAYERS=4 (set in ScriptArgs.__post_init__) +_MODEL_REGISTRY = { + "Inkling": "inkling", + "Inkling-4layer": "inkling", + "Inkling-Small": "inkling-small", + "Inkling-Small-4layer": "inkling-small", +} + + +@dataclass +class ScriptArgs(U.ExecuteTrainConfig): + run_id: str = U.create_run_id() + model_name: Literal["Inkling", "Inkling-4layer", "Inkling-Small", "Inkling-Small-4layer"] = "Inkling" + + task: Literal["dapo_math", "geo3k"] = "dapo_math" + fully_async: bool = False + rollout_num_nodes: int = 4 + enable_eval: bool = False + num_rollout: int = 100 + rollout_batch_size: int = 32 + global_batch_size: int = 64 + + hf_checkpoint: str | None = None + torch_dist: str | None = None + torch_dist_local: str | None = None + model_dir: str = "/root/models" + data_dir: str = "/root/datasets" + save_dir: str | None = None + megatron_path: str = "/root/Megatron-LM" + + num_gpus_per_node: int = 4 + rollout_num_gpus_per_engine: int = 16 + lr: float | None = None + rollout_max_response_len: int = 4096 + sglang_context_length: int = 8192 + train_offload_disk_dir: str = "/tmp/train_offload" + colocate: bool = field(init=False) + actor_num_nodes: int = field(init=False) + actor_num_gpus_per_node: int = field(init=False) + + enable_r3: bool = True + + extra_args: str = "" + + def __post_init__(self): + if self.model_name.endswith("-4layer"): + os.environ["MODEL_ARGS_NUM_LAYERS"] = "4" + if self.hf_checkpoint is None: + self.hf_checkpoint = f"{self.model_dir}/{self.model_name}" + if self.torch_dist is None: + self.torch_dist = f"{self.model_dir}/{self.model_name}_torch_dist" + if self.torch_dist_local is None: + self.torch_dist_local = self.torch_dist + if self.lr is None: + self.lr = 1e-6 + if self.fully_async: + assert not self.enable_eval, "the fully-async rollout function does not serve eval" + assert 0 < self.rollout_num_nodes < self.num_nodes + self.colocate = False + self.actor_num_nodes = self.num_nodes - self.rollout_num_nodes + else: + self.colocate = True + self.actor_num_nodes = self.num_nodes + self.actor_num_gpus_per_node = self.num_gpus_per_node + + @property + def is_mm(self): + return self.task == "geo3k" + + +def _get_parallel_config(args: ScriptArgs) -> str: + """Return parallel config args for tested GPU configurations. + + Only includes configurations that have been verified to work. + Raises NotImplementedError for untested configurations. + """ + total_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + + if args.fully_async and args.model_name == "Inkling" and total_gpus == 32: + # Validated disaggregated profile: 8 training nodes x 4 GPUs (+ rollout nodes). + return ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 2 " + "--expert-model-parallel-size 16 " + "--expert-tensor-parallel-size 1 " + ) + + if args.model_name == "Inkling-Small" and total_gpus == 32: + # TP4 x PP8 -> DP1; 42 = 7x5 + 7; EP<=TP*DP=4 -> 64 experts per rank. + return ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 8 " + "--decoder-last-pipeline-num-layers 7 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + ) + + if args.model_name in ("Inkling-4layer", "Inkling-Small-4layer") and args.actor_num_nodes == 1: + return ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + ) + + if args.actor_num_gpus_per_node == 4: + if total_gpus == 64: + return ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 4 " + "--decoder-last-pipeline-num-layers 15 " + "--expert-model-parallel-size 16 " + "--expert-tensor-parallel-size 1 " + ) + if total_gpus == 48: + return ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 3 " + "--expert-model-parallel-size 16 " + "--expert-tensor-parallel-size 1 " + ) + + raise NotImplementedError( + f"No pre-set parallel config for {total_gpus} GPUs. " + f"Please specify your parallel config in `run_inkling._get_parallel_config`." + ) + + +def _train(args: ScriptArgs): + topology = ( + f"{args.actor_num_nodes} train + {args.rollout_num_nodes} rollout nodes (fully-async)" + if args.fully_async + else f"{args.num_nodes} nodes (colocate)" + ) + print(f"running {args.model_name} full/{args.task} on {topology} x {args.num_gpus_per_node} GPUs") + + ckpt_args = ( + f"--hf-checkpoint {args.hf_checkpoint} " + f"--load {args.torch_dist_local} " + "--model-name inkling " + "--megatron-to-hf-mode raw " + "--no-load-optim --no-load-rng --finetune " + ) + if args.save_dir is not None: + ckpt_args += f"--save {args.save_dir}/{args.run_id}/checkpoints --save-interval 10 " + + rollout_args = "" + if args.fully_async: + rollout_args += "--fully-async " "--pause-generation-mode in_place " + rollout_args += ( + "--input-key prompt " + "--label-key label " + "--rollout-shuffle " + "--rm-type math " + f"--num-rollout {args.num_rollout} " + f"--rollout-batch-size {args.rollout_batch_size} " + "--n-samples-per-prompt 8 " + f"--rollout-max-response-len {args.rollout_max_response_len} " + "--rollout-temperature 1 " + f"--global-batch-size {args.global_batch_size} " + "--balance-data " + ) + eval_args = "" + match args.task: + case "dapo_math": + rollout_args += ( + f"--prompt-data {args.data_dir}/dapo-math-17k/dapo-math-17k.jsonl " "--apply-chat-template " + ) + if args.enable_eval: + eval_args = ( + "--eval-interval 5 " + f"--eval-prompt-data aime25 {args.data_dir}/aime-2025/aime-2025.jsonl " + "--eval-input-key prompt --eval-label-key label " + "--n-samples-per-eval-prompt 1 --eval-temperature 1 " + f"--eval-max-response-len {args.rollout_max_response_len} " + ) + case "geo3k": + rollout_args += f"--prompt-data {args.data_dir}/geo3k/geo3k_train.jsonl " + if args.enable_eval: + eval_args = ( + "--eval-interval 5 " + f"--eval-prompt-data geo3kval {args.data_dir}/geo3k/geo3k_val.jsonl " + "--eval-input-key prompt --eval-label-key label " + "--n-samples-per-eval-prompt 1 --eval-temperature 1 " + f"--eval-max-response-len {args.rollout_max_response_len} " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--entropy-coef 0.0 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + "--eps-clip-c 3.0 " + "--use-tis " + ) + if args.enable_r3: + grpo_args += "--use-rollout-routing-replay " + + optimizer_args = ( + "--optimizer adam " + f"--lr {args.lr} " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--use-distributed-optimizer " + "--no-check-for-nan-in-loss-and-grad " + ) + # the validated fully-async recipe reduces grads in bf16; colocate keeps fp32 + optimizer_args += "--grad-reduce-in-bf16 " if args.fully_async else "--accumulate-allreduce-grads-in-fp32 " + + perf_args = _get_parallel_config(args) + perf_args += "--recompute-granularity full " "--recompute-method uniform " "--recompute-num-layers 1 " + + if args.fully_async: + # disaggregated: nothing to offload, but the optimizer state does not fit + # the training GPUs while the step runs, so stream it through NVMe + optimizer_args += ( + "--stream-optimizer-state-to-disk " + f"--offload-train-disk-dir {args.train_offload_disk_dir} " + "--offload-train-disk-chunk-mb 256 " + ) + sglang_args = ( + "--rollout-num-gpus-per-engine 16 " + "--sglang-mem-fraction-static 0.80 " + "--sglang-max-running-requests 128 " + "--sglang-max-total-tokens 655360 " + "--sglang-cuda-graph-max-bs 128 " + ) + else: + optimizer_args += "--offload-train-target disk " f"--offload-train-disk-dir {args.train_offload_disk_dir} " + sglang_args = ( + f"--rollout-num-gpus-per-engine {args.rollout_num_gpus_per_engine} " + "--sglang-mem-fraction-static 0.6 " + "--sglang-max-running-requests 64 " + "--sglang-max-total-tokens 327680 " + ) + perf_args += "--micro-batch-size 1 " + + sglang_args += ( + "--sglang-attention-backend fa4 " + "--sglang-moe-runner-backend triton " + "--sglang-mamba-scheduler-strategy extra_buffer " + "--sglang-enable-multimodal " + f"--sglang-context-length {args.sglang_context_length} " + "--sglang-disable-custom-all-reduce " + ) + + inkling_args = "" + if args.is_mm: + # The mm provider wires the frozen HF vision/audio towers; it overrides the + # text provider from the model sh (MODEL_ARGS precede train_args; last one wins). + inkling_args = "--custom-model-provider-path miles_plugins.models.inkling.model.inkling_mm_model_provider " + + misc_args = ( + "--transformer-impl transformer_engine " + "--bf16 " + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--attention-softmax-in-fp32 " + "--no-bias-dropout-fusion " + "--distributed-timeout-minutes 30 " + f"--actor-num-nodes {args.actor_num_nodes} " + f"--actor-num-gpus-per-node {args.actor_num_gpus_per_node} " + f"--num-gpus-per-node {args.num_gpus_per_node} " + ) + if args.fully_async: + misc_args += ( + f"--rollout-num-gpus {args.rollout_num_nodes * args.num_gpus_per_node} " + "--update-weight-transfer-mode broadcast " + ) + else: + misc_args += "--colocate " + + extra_env_vars = { + "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1", + "SGLANG_OPT_USE_INKLING_FUSED_AR_SCONV_NORM": "false", + "SGLANG_SKIP_SGL_KERNEL_VERSION_CHECK": "1", + "MILES_SGLANG_DUMMY_LOAD": "0", + "SGLANG_SERVER_ENGINE_ROLLOUT_RETURN_LOGPROB": "1", + "RAY_memory_monitor_refresh_ms": "0", + "NCCL_MNNVL_ENABLE": "1", + "NCCL_NVLS_ENABLE": "0", + "NCCL_RAS_ENABLE": "0", + } + if args.fully_async: + extra_env_vars["MILES_EXPERIMENTAL_ROLLOUT_REFACTOR"] = "1" + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{eval_args} " + f"{grpo_args} " + f"{optimizer_args} " + f"{perf_args} " + f"{inkling_args} " + f"{sglang_args} " + f"{misc_args} " + f"{U.get_default_wandb_args(__file__, run_id=args.run_id)} " + f"{args.extra_args} " + ) + + U.execute_train( + train_args=train_args, + config=args, + num_gpus_per_node=args.num_gpus_per_node, + megatron_model_type=_MODEL_REGISTRY[args.model_name], + train_script="train_async.py" if args.fully_async else "train.py", + extra_env_vars=extra_env_vars, + megatron_path=args.megatron_path, + ) + + +@app.command() +@U.dataclass_cli +def train(args: ScriptArgs): + """Run training. Assumes HF checkpoint / torch_dist are already staged.""" + _train(args) + + +def _prepare_cp(args: ScriptArgs): + U.rsync_simple(path_src=args.torch_dist, path_dst=args.torch_dist_local) + + +@app.command() +@U.dataclass_cli +def prepare_cp(args: ScriptArgs): + """Copy the shared torch_dist checkpoint to node-local NVMe (torch_dist_local).""" + _prepare_cp(args) + + +@app.command() +@U.dataclass_cli +def full_train(args: ScriptArgs): + if args.torch_dist_local != args.torch_dist: + _prepare_cp(args) + else: + print(f"[full_train] Skipping rsync: torch_dist_local == torch_dist ({args.torch_dist})") + _train(args) + + +if __name__ == "__main__": + app() diff --git a/tests/e2e/megatron/model_scripts/test_inkling_small_4layer_ci.py b/tests/e2e/megatron/model_scripts/test_inkling_small_4layer_ci.py new file mode 100644 index 00000000000..20a6b03943f --- /dev/null +++ b/tests/e2e/megatron/model_scripts/test_inkling_small_4layer_ci.py @@ -0,0 +1,69 @@ +import os + +from scripts.run_inkling import _MODEL_REGISTRY, ScriptArgs, _train +from tests.ci.ci_register import register_cuda_ci +from tests.ci.metric_history import register_ci_gate + +import miles.utils.external_utils.command_utils as U + + +register_cuda_ci( + est_time=1800, + suite="stage-c-4-gpu-h200", + labels=["megatron", "model-scripts"], +) + +register_ci_gate(metric_key="train/grad_norm") +register_ci_gate(metric_key="train/ppo_kl") +register_ci_gate(metric_key="train/train_rollout_logprob_abs_diff") +register_ci_gate(metric_key="train/train_rollout_kl") +register_ci_gate(metric_key="rollout/raw_reward") + +_MODEL_ORG = "CharyZeng" + + +def _args() -> ScriptArgs: + return ScriptArgs( + model_name="Inkling-Small-4layer", + task="dapo_math", + num_nodes=1, + num_gpus_per_node=4, + rollout_num_gpus_per_engine=4, + num_rollout=2, + rollout_max_response_len=512, + sglang_context_length=1024, + extra_args=( + "--ci-test " + "--ci-disable-kl-checker " + "--check-weight-update-skip-list visual. audio. " + "--ci-disable-logprobs-checker " + "--disable-weights-backuper " + "--offload-train-target cpu " + ), + ) + + +def prepare(args: ScriptArgs): + U.exec_command(f"mkdir -p {args.model_dir} {args.data_dir}") + U.exec_command(f"hf download {_MODEL_ORG}/{args.model_name} --local-dir {args.hf_checkpoint}") + U.hf_download_dataset("zhuzilin/dapo-math-17k", data_dir=args.data_dir) + U.convert_checkpoint( + model_name=args.model_name, + megatron_model_type=_MODEL_REGISTRY[args.model_name], + num_gpus_per_node=args.num_gpus_per_node, + dir_dst=args.model_dir, + hf_checkpoint=args.hf_checkpoint, + megatron_path=args.megatron_path, + ) + + +def execute(args: ScriptArgs): + _train(args) + + +if __name__ == "__main__": + args = _args() + prepare(args) + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute(args) diff --git a/tools/convert_hf_to_torch_dist.py b/tools/convert_hf_to_torch_dist.py index c61a36dd249..285c6c1541e 100644 --- a/tools/convert_hf_to_torch_dist.py +++ b/tools/convert_hf_to_torch_dist.py @@ -27,6 +27,18 @@ def add_conversion_args(parser): default="raw", help="The method to convert megatron weights to hugging face weights for SGLang.", ) + parser.add_argument( + "--custom-model-provider-path", + type=str, + default=None, + help=( + "Path to a custom model provider function (e.g. for models like Inkling whose mcore " + "module structure differs from a plain GPTModel -- model-level embed_norm, custom " + "router/shared-experts). When set, the offline mcore model is built by this provider " + "(via miles' get_model_provider_func), then the mbridge bridge populates its weights. " + "Signature: def provider(pre_process, post_process, vp_stage=None) -> GPTModel." + ), + ) try: parser.add_argument("--padded-vocab-size", type=int, default=None) except Exception: @@ -55,7 +67,7 @@ def get_args(): def ceildiv(a, b): return -(a // -b) - if args.pipeline_model_parallel_size == 1 and world_size > 1: + if args.pipeline_model_parallel_size == 1 and world_size > 1 and not os.environ.get("CONVERT_KEEP_PP1"): pp_size = world_size while True: args.pipeline_model_parallel_size = pp_size