diff --git a/README.md b/README.md index 7242c6b..59eec0c 100644 --- a/README.md +++ b/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: diff --git a/tools/convert_to_hf.py b/tools/convert_to_hf.py index 02d9bc5..24032d8 100644 --- a/tools/convert_to_hf.py +++ b/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 diff --git a/tools/convert_to_torchspec.py b/tools/convert_to_torchspec.py new file mode 100755 index 0000000..21809ef --- /dev/null +++ b/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 \ + --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() diff --git a/torchspec/config/mooncake_config.py b/torchspec/config/mooncake_config.py index 6a258aa..c72be8f 100644 --- a/torchspec/config/mooncake_config.py +++ b/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", ""), diff --git a/torchspec/models/draft/keymap.py b/torchspec/models/draft/keymap.py new file mode 100644 index 0000000..3dabbf5 --- /dev/null +++ b/torchspec/models/draft/keymap.py @@ -0,0 +1,43 @@ +# 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. + +""" +Mapping TorchSpec weight keys to HF / vLLM / SGLang serving names. +""" + +DRAFT_WEIGHT_KEY_REMAP = [ + ("midlayer.", "layers.0."), + ("context_proj.", "fc."), + ("context_norm.", "hidden_norm."), + ("final_norm.", "norm."), +] + + +def to_export_keys(tensors: dict) -> dict: + """Rename internal draft keys to HF/SGLang export names (forward direction).""" + remapped = {} + for k, v in tensors.items(): + new_key = k + for internal_prefix, export_prefix in DRAFT_WEIGHT_KEY_REMAP: + if k.startswith(internal_prefix): + new_key = export_prefix + k[len(internal_prefix) :] + break + remapped[new_key] = v + return remapped