diff --git a/python/sglang/srt/debug_utils/model_truncator.py b/python/sglang/srt/debug_utils/model_truncator.py new file mode 100644 index 000000000..d2b6f4b37 --- /dev/null +++ b/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())