inference/README.md

Step-5 minimal inference

中文版本：README.zh.md

This directory contains the standalone Step-5 inference path. Full-model text and image-text generation have been validated with FP8 on one 8 x H200 node at TP=EP=8, and BF16 on two 8 x H200 nodes at TP=EP=16. The vision encoder uses PyTorch and BF16 weights. All commands below run from this directory; the inference files can be used independently.

The default sliding-attention path is PyTorch SDPA. QKNorm/RoPE, DSA sparse attention, block-scaled FP8 expert GEMM, and the critical MoE primitives use Triton/CUDA; model orchestration and NCCL collectives remain PyTorch.

The same ranks provide tensor parallelism for the dense backbone and expert parallelism for the 352 routed experts. TP=EP=8 places 44 experts on each rank; TP=EP=16 places 22. TP and EP share a process group; there is no independent TP × EP process grid.

Contents

File	What it is
convert.py	Converts an original Step-5 checkpoint into per-rank TP shards: tpN/model-r{rank}.safetensors + matching config-r{rank}.json, plus a shared tokenizer_files/.
generate.py	Greedy text or image-text inference. Loads one shard per torchrun rank; accepts a text --prompt, or exact prompt IDs with optional prepared --image-inputs.
model.py	Text/VL configuration, decoder layers, TP/EP execution, and image-feature injection into packed prefill. Imports compute kernels from kernel.py.
vision_encoder.py	PyTorch perception encoder with 2D RoPE, absolute positions and convolutional downsamplers.
kernel.py	The Triton/CUDA compute kernels: the sparse-attention (DSA) indexer + decode/prefill, block-scaled FP8 expert GEMM, the MoE primitives, and fused QKNorm/RoPE.
test_vision.py	Configuration, checkpoint-loading and multimodal inference regressions.
requirements.txt	Compatible runtime ranges: torch >=2.10,<2.14, transformers >=4.57.6,<6, safetensors >=0.7,<1, and triton >=3.6,<3.8.

Typical flow: convert.py (once, to shard a checkpoint) → generate.py (to run it). generate.py loads the model implemented by model.py, kernel.py and vision_encoder.py.

Expert storage follows the source configuration: absent or empty quantization_config selects BF16; quant_method="fp8" selects FP8 with BF16 exclusions from modules_to_not_convert. Conversion preserves weight dtypes.

Install

Use a CUDA-enabled PyTorch build matching the host driver. FP8 checkpoint inference at TP=8 requires eight visible GPUs with FP8 E4M3 support. The validated vision runtime uses PyTorch 2.13.0+cu129, Triton 3.7.1, Transformers 5.15.0, and Safetensors 0.8.0. The compatible version ranges are listed in requirements.txt.

Create a virtual environment and install the dependencies:

python3 -m venv --system-site-packages .venv
source .venv/bin/activate
python -m pip install -r requirements.txt

If the host does not already provide a compatible CUDA-enabled PyTorch, install it from the CUDA wheel index matching the driver before installing the remaining requirements:

python -m pip install \
  --index-url "<matching-pytorch-cuda-wheel-index>" \
  "torch>=2.10,<2.14"
python -m pip install -r requirements.txt

The release loads the fast tokenizer serialized in local tokenizer.json directly. Automatic tokenizer class selection in Transformers 5.15 is incompatible with these Step-5 assets and can collapse long prompts. Direct loading preserves the byte-level BPE input IDs and output decoding validated with Transformers 4.57.6.

Long-output stability

Generation uses the following long-output execution contract:

• TP partial results are added in rank order with FP32 accumulation.
• Sliding-window token-by-token decode uses PyTorch's default SDPA backend selection. This keeps the production path on the backend supported by the installed PyTorch/CUDA runtime instead of forcing the slower math backend.
• A request that emits EOS is removed from later decode steps and its cache row is compacted.

The checks cover row-position stability inside one active batch. They do not promise bitwise-identical output across different batch sizes, schedulers, GPU models, or software versions.

Checkpoints and launch modes

Run all commands from this directory. Choose the FP8 or BF16 section for your source checkpoint and use a separate output directory for each format. Conversion partitions weights and preserves their dtype; it does not quantize a BF16 checkpoint into FP8.

The source model directory supplies the authoritative config.json. convert.py writes tpN/model-r{rank}.safetensors and matching config-r{rank}.json files, with gqa-provider-shared-tp-v2 and expert_parallel_size=N, plus shared tokenizer assets in tokenizer_files/. generate.py reads these configs from --tp-dir and discovers the copied tokenizer automatically. The inference directory does not supply a model config. Use the weights and configs from the same conversion. If a config omits expert_parallel_size, the runtime uses tp_size.

FP8: one node with 8 GPUs

The FP8 path has been validated on 8 x H200, with TP=EP=8 and 44 routed experts per rank. The source config must declare quant_method="fp8"; excluded expert layers and the vision tower retain their BF16 weights.

Set the source and output directories:

export STEP5_FP8_CHECKPOINT="<fp8-checkpoint-directory>"
export STEP5_FP8_OUT="<fp8-shard-output-directory>"

Check checkpoint keys and shard shapes, then convert once to a fresh output directory. --check-only reads headers without loading weights or writing shards:

python3 convert.py \
  --checkpoint "$STEP5_FP8_CHECKPOINT" \
  --tp-size 8 --ep-size 8 --check-only

python3 convert.py \
  --checkpoint "$STEP5_FP8_CHECKPOINT" \
  --out-dir "$STEP5_FP8_OUT" \
  --tp-size 8 --ep-size 8

Launch text generation on the node with eight visible GPUs:

torchrun --standalone --nproc-per-node=8 generate.py \
  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
  --prompt "请介绍一下张量并行。" \
  --max-new-tokens 128

For image-text generation, prepare the matching checkpoint's prompt IDs and image tensors according to Vision inputs, then run:

torchrun --standalone --nproc-per-node=8 generate.py \
  --tp-dir "$STEP5_FP8_OUT/tp8" --ep-size 8 \
  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
  --max-new-tokens 128

BF16: 16 GPUs across two nodes

The full 92-layer BF16 VL model needs approximately 144.41 GiB per rank at TP/EP=8 for parameters alone, exceeding the validated H200 node's 140.40 GiB per GPU. At TP=EP=16, loaded parameters occupy 74.93 GiB per rank, with 22 routed experts per rank. These figures include replicated vision weights and a sharded output head, exclude MTP, and leave KV caches, RoPE, activations and communication as additional memory requirements.

The complete BF16 decoder, vision encoder and projector passed two-node text and single-image greedy generation on 16 H200 GPUs. All ranks produced identical tokens and reached EOS, with finite logits throughout. The measured peak PyTorch allocation was 76.37 GiB per GPU for batch size 1 and a 943-token image prompt, including 898 visual tokens. This excludes allocations outside PyTorch, such as NCCL and the CUDA context. Communication used NCCL TCP Socket; RDMA, long-context workloads and throughput have not been validated.

The full-model run streamed and sliced the original weights using the converter's TP rules and executed the released model forward path. The complete BF16 convert.py → generate.py --tp-dir file-based workflow below has not been separately run end to end; checkpoint key/shape checks and small-model BF16 conversion/loading/generation tests have passed.

Use two nodes with eight visible H200 GPUs each. Both nodes need the same code and compatible runtime, NCCL connectivity, and access to the same converted weights and tokenizer files at the configured paths. Image inputs must also be available with identical contents on both nodes.

Set these paths on both nodes. The source checkpoint must use BF16 weights and have no quantization config, or an empty one:

export STEP5_BF16_CHECKPOINT="<bf16-checkpoint-directory>"
export STEP5_BF16_OUT="<bf16-shard-output-directory>"

On one node only, check the source and convert once into 16 complete shards. Make the resulting tp16/ and tokenizer_files/ directories accessible to both nodes before launching:

python3 convert.py \
  --checkpoint "$STEP5_BF16_CHECKPOINT" \
  --tp-size 16 --ep-size 16 --check-only

python3 convert.py \
  --checkpoint "$STEP5_BF16_CHECKPOINT" \
  --out-dir "$STEP5_BF16_OUT" \
  --tp-size 16 --ep-size 16

On both nodes, set the same reachable node-0 address and free rendezvous port:

export MASTER_ADDR="<node-0-reachable-IP-or-hostname>"
export MASTER_PORT=29500

On node 0:

export NODE_RANK=0

On node 1:

export NODE_RANK=1

Run the following command on both nodes, with their respective NODE_RANK, for the same launch. It starts eight workers per node and 16 workers in total:

torchrun --nnodes=2 --nproc-per-node=8 \
  --node-rank="$NODE_RANK" \
  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
  generate.py \
  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
  --prompt "请介绍一下张量并行。" \
  --max-new-tokens 128

Alternatively, for image-text generation, prepare the inputs according to Vision inputs and run this command on both nodes:

torchrun --nnodes=2 --nproc-per-node=8 \
  --node-rank="$NODE_RANK" \
  --master-addr="$MASTER_ADDR" --master-port="$MASTER_PORT" \
  generate.py \
  --tp-dir "$STEP5_BF16_OUT/tp16" --ep-size 16 \
  --prompt-json prompt-ids.json --image-inputs image-inputs.safetensors \
  --max-new-tokens 128

The launcher sets WORLD_SIZE=2 x 8=16; shard configs and --ep-size must both use 16. --standalone is only used by the single-node FP8 example.

For exact text token IDs in either launch mode, replace --prompt and its value with --prompt-json <prompt-json-file>. The file accepts a JSON list or {"prompt_ids": [...]}. --prompt-json-batch accepts a batch of requests.

Limitations

• Full-model generation has been validated with FP8 on one 8 x H200 node and BF16 on two 8 x H200 nodes using streamed weights. The full BF16 file-based conversion/loading workflow remains unvalidated end to end.
• --ep-size defaults to 8 and currently must equal TP/world size; independent TP and EP process groups such as TP=4, EP=2 are not implemented. generate.py reads TP from the shard config (and checks it against WORLD_SIZE); it intentionally has no separate --tp-size switch.
• Sliding attention uses PyTorch SDPA.
• Generation is greedy. Continuous dynamic batching, paged KV cache, prefix caching and CUDA graphs are not provided. BF16 cross-node validation covers short text and single-image requests over NCCL TCP Socket.
• MTP layers are not loaded, so speculative decoding is unavailable.
• The checkpoint's factory FP8/BF16 expert layout is consumed directly; the converter and generated TP shards are checkpoint-layout specific.

Vision checkpoints (step5_vl)

Step5ForConditionalGeneration uses Step5VLConfig with top-level model_type=step5_vl, text_config.model_type=step5, and vision_config.model_type=perception_encoder. Text-only checkpoints use Step5Config and Step5ForCausalLM.

The Step-5 perception encoder uses PyTorch operations, non-causal SDPA, 2D RoPE, absolute position embeddings, LayerScale and two convolutional downsamplers. Projected image features enter the packed text decoder during prefill; text decoding uses the DSA cache and co-located TP/EP execution.

The validated VL checkpoint has the following geometry:

Component	Configuration
Text decoder	92 layers, hidden size 4096; 3 MTP layers excluded from ordinary generation
Vision Transformer	47 layers, width 1536, 16 attention heads, MLP size 8960
Patch embedding	RGB, 14×14 kernel and stride 14
Activation and normalization	QuickGELU, pre-LN, no post-LN, LayerScale 0.1
Position encoding	Absolute positions and 2D RoPE; no CLS token
Downsamplers	Two 3×3 convolutions, stride 2, padding 1; 1536→3072→6144
Projector	6144→4096, no bias
Image placeholder	Token ID 128001
Vision weights	667 BF16 tensors; approximately 3.96 GB per TP rank

Checkpoint keys are model.*, lm_head.*, vision_model.* and vit_large_projector.*. Vision and projector tensors are replicated on each TP rank. The checkpoint index determines the files and dtypes to load. The converter excludes MTP tensors at indices >= text_config.num_hidden_layers.

Use the FP8 or BF16 checkpoint-check and conversion commands above for VL checkpoints as well. Complete shards include the nested config, Step5ForConditionalGeneration identity, text, vision and projector weights.

Both resumed conversion and generation verify each shard's complete key set, shapes and TP/EP/rank metadata before reading weight data. Generation also rejects any missing or unexpected keys reported by the Hugging Face loader; incomplete vision weights cannot fall back to random initialization.

Vision inputs

The model API accepts preprocessed tensors:

Argument	Meaning
pixel_values	Overview images, [images, 3, H, W]; additional leading batch dimensions are flattened.
patch_pixel_values	Optional crops in image order, with the same layout.
num_patches	One nonnegative crop count per overview image. Defaults to zero crops.
image_embeds	Alternative already projected features, [..., text_hidden_size]; excludes all pixel/crop arguments.

The released geometry produces 169 tokens for each 728×728 overview and 81 for each 504×504 crop. Each image contributes its crops first, then its overview. Concatenate images in prompt order across packed requests. The number of image_token_id placeholders must equal the total feature-token count; <im_start> / <im_end> remain ordinary text embeddings. Image encoding runs only during prefill. Decode uses the existing cache.

Image resizing, normalization, crop selection and prompt construction must come from the matching checkpoint processor. This standalone path consumes those prepared results; it does not define a new PIL/JPEG processing policy. Save a prepared batch with safetensors.torch.save_file, for example:

import torch
from safetensors.torch import save_file

save_file({
    "pixel_values": pixel_values.contiguous(),
    "patch_pixel_values": patch_pixel_values.contiguous(),
    "num_patches": torch.tensor(num_patches, dtype=torch.int64),
}, "image-inputs.safetensors")

Use the image-text launch command in the FP8 or BF16 section with prompt-ids.json (including image placeholders) and image-inputs.safetensors.

--prompt-json-batch supports a mixture of text and image requests. Every rank reads the same prepared image file. The Python API accepts these arguments in Step5ForConditionalGeneration.forward and generate_greedy as well.

Run the multimodal integration regression with the optional test dependency:

python -m pip install pytest
python -m pytest -q test_vision.py

vision_select_layer=-1 and language-frequency 2D RoPE are supported. Other layer-selection or frequency modes fail explicitly.

Validation results

The following results were obtained on H200 with PyTorch 2.13.0+cu129, Triton 3.7.1 and Transformers 5.15.0:

Check	Result
Vision integration, checkpoint identity and BF16 loading tests	20 passed
Configuration, generation, serving, evaluation and cache regressions	203 passed, 93 subtests passed
GPU operator gates	All 10 passed
Batch and NCCL TP stability	Passed
FP8 checkpoint schema, TP8	2653 runtime tensors validated; 51 MTP tensors excluded
BF16 checkpoint schema, TP8 and TP16	2398 runtime tensors validated; 51 MTP tensors excluded
BF16 small-model generation	TP1/TP8 pixel and projected-embedding inputs produced identical output token IDs
Full vision numerical validation	Overview and crop outputs matched exactly; maximum absolute error 0
TP8 mixed-precision fixture	Pixel and projected-embedding inputs produced identical output token IDs
HF file-shard loading, complete real rank 0	85.4 GB shard loaded; real-image features [898,4096] matched exactly
Full real FP8 checkpoint TP8 image generation, streamed weights	Correctly described a white Samoyed on grass; 27 generated tokens, normal EOS, identical output on all ranks
Full real BF16 checkpoint, two-node TP16/EP16, streamed weights	Text and single-image generation passed; finite logits, normal EOS and identical tokens on all 16 ranks; 74.93 GiB parameters/rank, 76.37 GiB peak PyTorch allocation/GPU

The image case used one 728×728 overview and nine 504×504 crops, totaling 898 visual tokens. Image encoding runs in PyTorch with fused GPU attention. These results cover functional and numerical correctness for the tested inputs; they do not constitute a multimodal benchmark or OCR/video quality evaluation.