Add script to create a model with fewer layers for debugging (#13284)

python/sglang/srt/debug_utils/model_truncator.py

@@ -0,0 +1,112 @@
+# This file also references Slime :: fp8_cast_bf16.py
+import json
+import os
+import re
+from argparse import ArgumentParser
+from pathlib import Path
+from typing import Dict
+
+import torch
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file, save_file
+
+
+def main(args):
+    dir_input = Path(_maybe_snapshot_download(args.input))
+    dir_output = Path(args.output)
+    print(f"{dir_input=} {dir_output=}")
+
+    dir_output.mkdir(parents=True, exist_ok=True)
+
+    for pattern in ["generation_config.json", "*.py", "tokenizer*"]:
+        os.system(f"cp -rf {dir_input}/{pattern} {dir_output}")
+
+    _transform_json(
+        dir_input,
+        dir_output,
+        "config.json",
+        lambda data: _transform_config(args, data),
+    )
+
+    safetensors_index = _transform_json(
+        dir_input,
+        dir_output,
+        "model.safetensors.index.json",
+        lambda data: _transform_safetensors_index(args, data),
+    )
+
+    for path_input_safetensors in sorted(list(dir_input.glob("*.safetensors"))):
+        path_output_safetensors = dir_output / path_input_safetensors.relative_to(
+            dir_input
+        )
+
+        state_dict = load_file(path_input_safetensors)
+        _transform_safetensors_file(
+            state_dict, safetensors_index, debug_name=str(path_output_safetensors)
+        )
+        if len(state_dict) > 0:
+            print(f"Save {len(state_dict)} tensors to {path_output_safetensors}")
+            save_file(state_dict, path_output_safetensors)
+        else:
+            print(f"Skip saving {path_output_safetensors} since it is empty")
+
+
+def _maybe_snapshot_download(path):
+    if Path(path).exists():
+        return path
+    return snapshot_download(path)
+
+
+def _transform_json(dir_input, dir_output, filename, fn):
+    data = json.loads((dir_input / filename).read_text())
+    fn(data)
+    (dir_output / filename).write_text(json.dumps(data, indent=4))
+    return data
+
+
+def _transform_config(args, config_json):
+    config_json["num_hidden_layers"] = args.keep_num_layers
+
+
+def _transform_safetensors_index(args, safetensors_index):
+    weight_map = safetensors_index["weight_map"]
+    weight_map = {
+        name: loc for name, loc in weight_map.items() if _filter_tensor_name(args, name)
+    }
+    safetensors_index["weight_map"] = weight_map
+
+
+def _transform_safetensors_file(
+    state_dict: Dict[str, torch.Tensor], safetensors_index, debug_name: str
+):
+    names_to_remove = set(state_dict) - set(safetensors_index["weight_map"])
+    print(f"Remove {list(names_to_remove)} in {debug_name}")
+    for name in names_to_remove:
+        del state_dict[name]
+
+
+def _filter_tensor_name(args, tensor_name: str):
+    # We focus on DeepSeek-like names currently, but can be easily extended to more kinds of models
+    m = re.match(r"^model.layers.(\d+).*", tensor_name)
+    if m is None:
+        return True
+
+    layer_id = int(m.group(1))
+    return layer_id < args.keep_num_layers
+
+
+if __name__ == "__main__":
+    """
+    Example:
+    python -m sglang.srt.debug_utils.model_truncator --input deepseek-ai/DeepSeek-V3-0324 --output /tmp/DeepSeek-V3-0324-5layer
+    hf upload my_name/DeepSeek-V3-0324-5layer /tmp/DeepSeek-V3-0324-5layer
+
+    Alternatively, the following may be used on-the-fly.
+    But this may not be useful to test RL frameworks, and sometimes it may have issues.
+        --json-model-override-args '{"num_hidden_layers": 5}'
+    """
+    parser = ArgumentParser(description="Create truncated model for fast debugging.")
+    parser.add_argument("--input", type=str, required=True)
+    parser.add_argument("--output", type=str, required=True)
+    parser.add_argument("--keep-num-layers", type=int, default=5)
+    main(parser.parse_args())