Skip to content

feat(tools): add HF model normalization utility - #2073

Open
caixianzhang wants to merge 4 commits into
InternLM:mainfrom
caixianzhang:feature/model-normalize
Open

feat(tools): add HF model normalization utility#2073
caixianzhang wants to merge 4 commits into
InternLM:mainfrom
caixianzhang:feature/model-normalize

Conversation

@caixianzhang

@caixianzhang caixianzhang commented Sep 6, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • add an XTuner xtuner.tools.model_normalize CLI for HF repack and FP8 conversion
  • preserve MTP tensors/configuration and copy tokenizer/chat template/generation config assets
  • add bounded FP8 save workers, reference-guided or explicit heuristic policy, and one-click shell wrappers
  • add usage documentation and focused CLI/MTP tests

Scope

This PR intentionally provides conversion/repacking only. It does not run full-model validation, SHA256 scans, inference services, or Hub uploads.

Validation

  • ruff check xtuner/tools/model_normalize tests/tools/test_model_normalize_cli.py
  • PYTHONPATH=. pytest -q --confcutdir=tests/tools tests/tools/test_model_normalize_cli.py (2 passed, 1 skipped when safetensors is unavailable)
  • bash xtuner/tools/model_normalize/run_model_normalize.sh --help

run_glm52.sh usage

The GLM-5.2 example accepts one positional argument:

$1: bf16 | fp8

It does not define other positional arguments; extra arguments are ignored.

Environment variables:

Variable BF16 FP8 Default Description
SOURCE_DIR required required Input HF model directory
OUTPUT_ROOT required required Output root directory
SHARD_SIZE_GB optional optional 4 Target shard size in GB
REFERENCE_DIR not used required FP8 reference model directory
MAX_SAVE_WORKERS not used optional 4 Number of parallel FP8 shard-save workers

BF16 example:

SOURCE_DIR=/path/to/glm52-hf-source \
OUTPUT_ROOT=/path/to/release \
SHARD_SIZE_GB=4 \
bash xtuner/tools/model_normalize/examples/run_glm52.sh bf16

FP8 example:

SOURCE_DIR=/path/to/glm52-hf-source \
OUTPUT_ROOT=/path/to/release \
REFERENCE_DIR=/path/to/glm52-fp8-reference \
SHARD_SIZE_GB=4 \
MAX_SAVE_WORKERS=4 \
bash xtuner/tools/model_normalize/examples/run_glm52.sh fp8

The output directories are fixed relative to OUTPUT_ROOT:

BF16: ${OUTPUT_ROOT}/20_hf_bf16_mtp
FP8:  ${OUTPUT_ROOT}/20_hf_fp8_mtp

Real-model smoke test

The wrapper was also run against the GLM-5.2-derived HF checkpoint in the
glm5-2-8gpu-99022910-cfcd2 Job. Because the Job has no external network
access, the PR code was staged through the shared storage into a temporary
directory; the existing repository checkout was not modified.

Input:

/mnt/shared-storage-gpfs2/intern-pretrain-shared02/zhouyunhua_cp/Delivery/muonsplit-and-clipsplit/exp_agent_0829a_muon_0_1/20260829144456/hf-425

Commands used:

SOURCE_DIR=/mnt/shared-storage-gpfs2/intern-pretrain-shared02/zhouyunhua_cp/Delivery/muonsplit-and-clipsplit/exp_agent_0829a_muon_0_1/20260829144456/hf-425 \
OUTPUT_ROOT=/mnt/shared-storage-user/llmit1/user/zhangcaixian/my_xtunner/run_glm52_fp8_20260906 \
REFERENCE_DIR=/mnt/shared-storage-user/llmit1/user/zhangcaixian/my_xtunner/release_artifacts/glm5p2-0829a-hf425/20_hf_fp8_mtp \
bash xtuner/tools/model_normalize/examples/run_glm52.sh fp8
SOURCE_DIR=/mnt/shared-storage-gpfs2/intern-pretrain-shared02/zhouyunhua_cp/Delivery/muonsplit-and-clipsplit/exp_agent_0829a_muon_0_1/20260829144456/hf-425 \
OUTPUT_ROOT=/mnt/shared-storage-user/llmit1/user/zhangcaixian/my_xtunner/run_glm52_bf16_20260906 \
bash xtuner/tools/model_normalize/examples/run_glm52.sh bf16

Results:

  • FP8: 2443/2443 input shards processed, 177/177 output shards written, about 704 GB.
  • BF16: 353/353 output shards written, about 1.4 TB.
  • Both outputs contain model.safetensors.index.json and config.json; both
    conversion sessions exited normally.
  • This smoke test confirms conversion completion and basic artifact presence;
    it does not replace full-model loading, semantic weight validation, or
    inference testing. generation_config.json was absent because it was not
    present in the input checkpoint.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants