fix: Mooncake RDMA memory for aux_layers > 3

README.md

@@ -191,6 +191,25 @@ training:
 output_dir: /path/to/new_run
 ```
 
+Start a new run from an existing Hugging Face checkpoint. `--input` accepts a Hugging Face Hub repo id (downloaded automatically), a local HF/safetensors directory or `.safetensors` file, or a TorchSpec DCP checkpoint dir:
+
+```bash
+python tools/convert_to_torchspec.py \
+    --input org/dflash-checkpoint \
+    --config torchspec/config/dspark_draft_config_qwen36_35b.json \
+    --output ./outputs/dspark_init
+```
+
+Then warm-start from the generated init:
+
+```yaml
+training:
+  load_path: ./outputs/dspark_init
+  continual_training: true
+```
+
+Using this technique, you can also warm-start DSpark training runs from a pre-trained DFlash model for faster convergence.
+
 ## Checkpoint Conversion
 
 Convert an FSDP checkpoint to HuggingFace format:
@@ -229,10 +248,14 @@ Set `TORCHSPEC_LOG_LEVEL=DEBUG` for more verbose logging when diagnosing issues:
 TORCHSPEC_LOG_LEVEL=DEBUG ./examples/qwen3-8b-single-node/run.sh
 ```
 
-## Mooncake SEGFAULT
+### Mooncake SEGFAULT
 
 Current Mooncake version has a bug with TCP-only hosts causing a SEGFAULT error. Set `MC_STORE_MEMCPY=0` until the [upstream issue](https://github.com/kvcache-ai/Mooncake/issues/1986) is fixed.
 
+### RDMA failure
+
+If you get the error `... libcudart symbols not found globally. Make sure PyTorch with CUDA is installed before using TileLang`, reinstall mooncake to match your system CUDA version, i.e. if you run on CUDA 13, mooncake is probably installed for cu12 wheels, re-install it using `uv pip uninstall mooncake-transfer-engine && uv pip install mooncake-transfer-engine-cuda13`.
+
 ### Per-Rank File Logging
 
 Set `TORCHSPEC_LOG_DIR` to an absolute path on a shared filesystem (NFS) to enable per-rank log files for every Ray actor on both training and inference:

tools/convert_to_hf.py

@@ -58,6 +58,7 @@
 from typing_extensions import override
 
 from torchspec.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel
+from torchspec.models.draft.keymap import DRAFT_WEIGHT_KEY_REMAP
 
 logging.basicConfig(
     level=logging.INFO,
@@ -127,19 +128,13 @@ def set_up_planner(
 # ── Export fixups ────────────────────────────────────────────────────────────
 
 # vLLM / SGLang supports a different naming convention for some weight keys.
-_WEIGHT_KEY_REMAP = [
-    ("midlayer.", "layers.0."),
-    ("context_proj.", "fc."),
-    ("context_norm.", "hidden_norm."),
-    ("final_norm.", "norm."),
-]
 
 
 def _remap_weight_keys(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
     remapped = {}
     for k, v in tensors.items():
         new_key = k
-        for old_prefix, new_prefix in _WEIGHT_KEY_REMAP:
+        for old_prefix, new_prefix in DRAFT_WEIGHT_KEY_REMAP:
             if k.startswith(old_prefix):
                 new_key = new_prefix + k[len(old_prefix) :]
                 break

tools/convert_to_torchspec.py

@@ -0,0 +1,221 @@
+#!/usr/bin/env python
+# Copyright (c) 2026 LightSeek Foundation
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+"""
+Build a warm-start init checkpoint for a TorchSpec draft model from an external
+draft checkpoint, such as HF weights.
+
+Source (``--input``) may be:
+  * a Hugging Face Hub repo id (downloaded automatically), or
+  * a local HF / safetensors dir (single or sharded) or a ``.safetensors`` file, or
+  * a TorchSpec torch-DCP dir (another run's checkpoint or an ``iter_XXXX`` dir).
+
+This is the inverse of ``tools/convert_to_hf.py``.
+
+Usage:
+  python tools/convert_to_torchspec.py \
+      --input  <foreign ckpt dir> \
+      --config torchspec/config/dspark_draft_config_qwen36_35b.json \
+      --output outputs/dspark_zai_warmstart_init
+
+Then in the training yaml:
+  training:
+    load_path: outputs/dspark_zai_warmstart_init
+    continual_training: true
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+import torch.nn as nn
+
+# Reuse the proven DCP flat-loader helpers from the sibling exporter.
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from convert_to_hf import _EmptyStateDictLoadPlanner, _WrappedStorageReader  # noqa: E402
+
+from torchspec.models.draft import AutoDraftModelConfig, AutoEagle3DraftModel  # noqa: E402
+from torchspec.models.draft.keymap import DRAFT_WEIGHT_KEY_REMAP  # noqa: E402
+from torchspec.training.checkpoint import ModelState  # noqa: E402
+
+
+def _load_dcp_flat(model_dir: str) -> dict:
+    sd: dict = {}
+    dcp.state_dict_loader._load_state_dict(
+        sd,
+        storage_reader=_WrappedStorageReader(model_dir),
+        planner=_EmptyStateDictLoadPlanner(),
+        no_dist=True,
+    )
+    return sd
+
+
+def _resolve_input(input_path: str) -> str:
+    """Return a local path for input_path, downloading it from the Hugging Face Hub
+    if it is a repo id rather than an existing local path."""
+    if os.path.exists(input_path):
+        return input_path
+    from huggingface_hub import snapshot_download
+
+    print(f"'{input_path}' is not a local path; resolving as a Hugging Face Hub repo id...")
+    return snapshot_download(repo_id=input_path, allow_patterns=["*.safetensors", "*.json"])
+
+
+def _load_safetensors(path: str):
+    """Load a single-file or sharded (index.json) safetensors checkpoint. Returns
+    (state_dict, source) or None if no safetensors are found at path."""
+    from safetensors.torch import load_file
+
+    if path.endswith(".safetensors") and os.path.isfile(path):
+        return load_file(path), f"safetensors:{path}"
+    single = os.path.join(path, "model.safetensors")
+    if os.path.exists(single):
+        return load_file(single), f"safetensors:{single}"
+    index = os.path.join(path, "model.safetensors.index.json")
+    if os.path.exists(index):
+        import json
+
+        with open(index) as f:
+            weight_map = json.load(f)["weight_map"]
+        shards = sorted(set(weight_map.values()))
+        sd: dict = {}
+        for shard in shards:
+            sd.update(load_file(os.path.join(path, shard)))
+        return sd, f"safetensors[{len(shards)} shards]:{path}"
+    return None
+
+
+def _load_flat(input_path: str):
+    """Return (flat_state_dict, source_description). Accepts a HF Hub repo id, a
+    local HF / safetensors dir or file, or a TorchSpec torch-DCP dir."""
+    path = _resolve_input(input_path)
+    # torch-DCP?
+    for cand in (path, os.path.join(path, "model")):
+        if os.path.exists(os.path.join(cand, ".metadata")):
+            return _load_dcp_flat(cand), f"dcp:{cand}"
+    tracker = os.path.join(path, "latest_checkpointed_iteration.txt")
+    if os.path.exists(tracker):
+        it = int(open(tracker).read().strip())
+        cand = os.path.join(path, f"iter_{it:07d}", "model")
+        if os.path.exists(os.path.join(cand, ".metadata")):
+            return _load_dcp_flat(cand), f"dcp:{cand}"
+    # safetensors (single / sharded / file)?
+    st = _load_safetensors(path)
+    if st is not None:
+        return st
+    raise FileNotFoundError(
+        f"No torch-DCP (.metadata) or safetensors found for input '{input_path}' (resolved: {path})"
+    )
+
+
+def _remap_to_model(raw: dict, model_keys: set) -> dict:
+    """Strip any wrapper/draft prefix and remap sglang export names toward the
+    target model's keys, applying an alias only when it lands on a real model key
+    (so ``fc.`` -> ``context_proj.`` fires while a valid ``layers.0.`` is kept)."""
+    out: dict = {}
+    for k, v in raw.items():
+        if not isinstance(v, torch.Tensor):
+            continue
+        dk = k.split("draft_model.")[-1] if "draft_model." in k else k
+        if dk in ("t2d", "d2t"):  # vocab-pruning maps, not weights
+            continue
+        target = dk
+        if target not in model_keys:
+            for internal_prefix, export_prefix in DRAFT_WEIGHT_KEY_REMAP:
+                if dk.startswith(export_prefix):
+                    cand = internal_prefix + dk[len(export_prefix) :]
+                    if cand in model_keys:
+                        target = cand
+                        break
+        out[target] = v
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    ap.add_argument(
+        "--input",
+        required=True,
+        help="Foreign draft checkpoint: a HF Hub repo id, a local HF/safetensors dir or "
+        ".safetensors file, or a TorchSpec DCP dir",
+    )
+    ap.add_argument("--config", required=True, help="Target draft model config.json (e.g. DSpark)")
+    ap.add_argument("--output", required=True, help="Output init dir (use as training.load_path)")
+    ap.add_argument("--dtype", default="bfloat16")
+    args = ap.parse_args()
+
+    dtype = getattr(torch, args.dtype)
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29655")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    dist.init_process_group("gloo", rank=0, world_size=1)
+
+    cfg = AutoDraftModelConfig.from_file(args.config)
+    model = AutoEagle3DraftModel.from_config(cfg).to(dtype)
+    model_keys = set(model.state_dict().keys())
+
+    raw, src = _load_flat(args.input)
+    remapped = _remap_to_model(raw, model_keys)
+    matched = {k: v.to(dtype) for k, v in remapped.items() if k in model_keys}
+    result = model.load_state_dict(matched, strict=False)
+
+    ignored = sorted(set(remapped) - model_keys)
+    print(f"source: {src}")
+    print(f"matched {len(matched)}/{len(model_keys)} target keys")
+    print(f"left at fresh init ({len(result.missing_keys)}): {sorted(result.missing_keys)}")
+    if ignored:
+        print(f"ignored source keys ({len(ignored)}): {ignored}")
+
+    if len(matched) == 0:
+        raise SystemExit(
+            "ERROR: no source key matched the target model — check formats / key names."
+        )
+    if len(matched) < 0.5 * len(model_keys):
+        raise SystemExit(
+            f"ERROR: only {len(matched)}/{len(model_keys)} keys matched — likely a key-mapping "
+            "failure; refusing to write a half-initialized init."
+        )
+
+    class Wrapper(nn.Module):
+        def __init__(self, draft_model):
+            super().__init__()
+            self.draft_model = draft_model
+
+    model_dir = os.path.join(args.output, "iter_0000001", "model")
+    os.makedirs(model_dir, exist_ok=True)
+    dcp.save({"model_state": ModelState(Wrapper(model))}, checkpoint_id=model_dir)
+    with open(os.path.join(args.output, "latest_checkpointed_iteration.txt"), "w") as f:
+        f.write("1")
+
+    print(f"\nWrote warm-start init to {args.output}")
+    print("Point the training yaml at it:")
+    print(f"  training:\n    load_path: {args.output}\n    continual_training: true")
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

torchspec/config/mooncake_config.py

@@ -63,6 +63,7 @@ class MooncakeConfig:
     get_batch_size: int = 1
     max_seq_len: int = 8192
     hidden_dim: int = 4096
+    num_aux_layers: int = 3
     async_put_pool_size: int | None = None
     store_full_error_codes: Tuple[int, ...] = (-200,)
     store_full_wait_seconds: float = 0.5
@@ -91,6 +92,7 @@ def __post_init__(self):
                 max_seq_len=self.max_seq_len,
                 batch_size=1,
                 hidden_dim=self.hidden_dim,
+                num_aux_layers=self.num_aux_layers,
                 safety_margin=2.0,
             )
 
@@ -104,8 +106,24 @@ def __post_init__(self):
                 max_seq_len=self.max_seq_len,
                 batch_size=self.get_batch_size,
                 hidden_dim=self.hidden_dim,
+                num_aux_layers=self.num_aux_layers,
             )
 
+    @staticmethod
+    def _infer_num_aux_layers(args) -> int:
+        draft_cfg = getattr(args, "draft_model_config_obj", None)
+        if draft_cfg is not None:
+            target_layer_ids = getattr(draft_cfg, "target_layer_ids", None)
+            if target_layer_ids:
+                return len(target_layer_ids)
+            num_target_layers = getattr(draft_cfg, "num_target_layers", None)
+            if num_target_layers:
+                return int(num_target_layers)
+        aux_layers = getattr(args, "aux_hidden_states_layers", None)
+        if aux_layers:
+            return len(aux_layers)
+        return 3
+
     @classmethod
     def from_flat_args(cls, args) -> "MooncakeConfig":
         """Create config from a flat args namespace (mooncake_* prefixed fields).
@@ -154,6 +172,7 @@ def from_flat_args(cls, args) -> "MooncakeConfig":
                 getattr(args, "max_seq_length", 8192),
             ),
             "hidden_dim": getattr(args, "mooncake_hidden_dim", 4096),
+            "num_aux_layers": cls._infer_num_aux_layers(args),
         }
 
         if metadata_server is not None:
@@ -174,6 +193,7 @@ def export_env(self) -> None:
         os.environ["MOONCAKE_GLOBAL_SEGMENT_SIZE"] = str(self.global_segment_size)
         os.environ["MOONCAKE_LOCAL_BUFFER_SIZE"] = str(self.local_buffer_size)
         os.environ["MOONCAKE_HOST_BUFFER_SIZE"] = str(self.host_buffer_size)
+        os.environ["MOONCAKE_NUM_AUX_LAYERS"] = str(self.num_aux_layers)
         os.environ["MOONCAKE_PROTOCOL"] = self.protocol
         os.environ["MOONCAKE_DEVICE_NAME"] = self.device_name
         os.environ["MOONCAKE_ENABLE_GPU_DIRECT"] = "1" if self.enable_gpu_direct else "0"
@@ -225,6 +245,8 @@ def from_env(cls) -> "MooncakeConfig":
         pool_size_env = os.getenv("MOONCAKE_ASYNC_PUT_POOL_SIZE")
         async_put_pool_size = int(pool_size_env) if pool_size_env is not None else None
 
+        num_aux_layers = int(os.getenv("MOONCAKE_NUM_AUX_LAYERS", "3"))
+
         return cls(
             local_hostname=os.getenv("MOONCAKE_LOCAL_HOSTNAME", "localhost"),
             metadata_server=os.getenv(
@@ -239,6 +261,7 @@ def from_env(cls) -> "MooncakeConfig":
             ),
             local_buffer_size=int(os.getenv("MOONCAKE_LOCAL_BUFFER_SIZE", str(512 * 1024 * 1024))),
             host_buffer_size=host_buffer_size,
+            num_aux_layers=num_aux_layers,
             async_put_pool_size=async_put_pool_size,
             protocol=os.getenv("MOONCAKE_PROTOCOL", "tcp"),
             device_name=os.getenv("MOONCAKE_DEVICE_NAME", ""),