LLaMA Streaming Token Cutoff After ZeRO‑3 Fine‑tune On 8‑GPU

Q: Root Cause Analysis

The cutoff originates from a mismatch between the model state dict sharding performed by ZeRO‑3 during training and the parameter loading logic used at inference time. ZeRO‑3 stores each tensor in a separate checkpoint shard (one per GPU) and also records a zero_stage flag in the checkpoint metadata. When the inference script loads the checkpoint with torch.load (or accelerate) without informing DeepSpeed that the checkpoint is ZeRO‑3, the loader attempts to reconstruct the full model on a single device. This triggers an out‑of‑memory (OOM) condition on the first GPU, causing the generation loop to abort silently after the first successful forward pass.

In this article

Table of Contents

Problem: Premature Termination of LLaMA Token Stream after ZeRO‑3 Fine‑Tune on 8‑GPU

After fine‑tuning Meta LLaMA‑7B with DeepSpeed ZeRO‑3 on an 8‑GPU NVIDIA A100 cluster, the streaming inference API stops emitting tokens after a few dozen generations. The client receives a truncated response and the server logs show no explicit error, only a silent exit from the generation loop.

Root Cause Analysis

The cutoff originates from a mismatch between the model state dict sharding performed by ZeRO‑3 during training and the parameter loading logic used at inference time. ZeRO‑3 stores each tensor in a separate checkpoint shard (one per GPU) and also records a zero_stage flag in the checkpoint metadata. When the inference script loads the checkpoint with torch.load (or accelerate) without informing DeepSpeed that the checkpoint is ZeRO‑3, the loader attempts to reconstruct the full model on a single device. This triggers an out‑of‑memory (OOM) condition on the first GPU, causing the generation loop to abort silently after the first successful forward pass.

Additional contributing factors:

DeepSpeed’s stage3 checkpoint includes optimizer_state and lr_scheduler_state that are unnecessary for inference but increase memory pressure.
The streaming API uses torch.no_grad() but does not wrap the model with deepspeed.init_inference, so tensor partitioning is not re‑applied.
The default max_new_tokens in the generation loop is set to None, relying on the model’s internal eos_token_id. When the model is partially loaded, the EOS token is never reached, and the loop exits after the internal safety limit (usually 256 tokens) is hit, appearing as a “cutoff”.

Investigation and Debugging

Inspect checkpoint metadata

python - <<'PY'
import json, torch
meta = torch.load('checkpoints/step_5000/meta.pt', map_location='cpu')
print(json.dumps(meta, indent=2))
PY

Typical output showing ZeRO‑3:

{
  "zero_stage": 3,
  "zero_version": "0.9.0",
  "model_sharding": true,
  "optimizer_sharding": true,
  "lr_scheduler_sharding": true
}

Reproduce the failure with a minimal script

#!/usr/bin/env python
import torch
from transformers import LlamaForCausalLM, LlamaTokenizer

tokenizer = LlamaTokenizer.from_pretrained('meta-llama/Llama-2-7b')
model = LlamaForCausalLM.from_pretrained(
    'checkpoints/step_5000',
    torch_dtype=torch.float16,
    device_map='auto'  # <-- triggers full model load
)

prompt = "Explain the difference between TCP and UDP."
input_ids = tokenizer(prompt, return_tensors='pt').input_ids.cuda()
with torch.no_grad():
    outputs = model.generate(input_ids, max_new_tokens=128, do_sample=False)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Running on the 8‑GPU node yields:

[2024-10-20 12:34:56] INFO  torch.cuda: Out of memory. Tried to allocate 4.00 GiB...
Traceback (most recent call last):
  ...
RuntimeError: CUDA out of memory.

However, the production streaming service swallows the exception and simply stops emitting tokens.

Check DeepSpeed inference initialization

import deepspeed

model = LlamaForCausalLM.from_pretrained('meta-llama/Llama-2-7b')
model = deepspeed.init_inference(
    model,
    mp_size=8,
    dtype=torch.float16,
    replace_method='auto',
    replace_with_kernel_inject=True
)

When this block is omitted, the model is loaded without ZeRO‑3 aware partitioning, reproducing the cutoff.

Resolution

The fix consists of two coordinated changes:

Export a ZeRO‑3‑compatible inference checkpoint

During training, add a post‑training step to consolidate the sharded checkpoint into a zero‑inference format:

# After training completes
deepspeed checkpoint --zero3_to_fp16 --output_dir inference_ckpt

This command strips optimizer and scheduler states and writes a single model.pt that can be loaded without DeepSpeed sharding.

Load the checkpoint with DeepSpeed inference API

Update the serving script to use deepspeed.init_inference and point to the new checkpoint directory:

#!/usr/bin/env python
import torch, deepspeed
from transformers import LlamaForCausalLM, LlamaTokenizer

tokenizer = LlamaTokenizer.from_pretrained('meta-llama/Llama-2-7b')
model = LlamaForCausalLM.from_pretrained(
    'inference_ckpt',
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

model = deepspeed.init_inference(
    model,
    mp_size=8,
    dtype=torch.float16,
    replace_method='auto',
    replace_with_kernel_inject=True
)

def stream_generate(prompt, max_new=256):
    input_ids = tokenizer(prompt, return_tensors='pt').input_ids.cuda()
    with torch.no_grad():
        for token in model.generate(
            input_ids,
            max_new_tokens=max_new,
            do_sample=False,
            output_scores=False,
            return_dict_in_generate=False,
        ):
            yield tokenizer.decode(token, skip_special_tokens=True)

# Example streaming usage
for chunk in stream_generate("Explain the difference between TCP and UDP."):
    print(chunk, end='', flush=True)

Key differences from the failing version:

Aspect	Before	After
Checkpoint format	ZeRO‑3 sharded (multiple .pt files)	Consolidated inference checkpoint (single model.pt)
Loading path	torch.load + device_map='auto'	deepspeed.init_inference with mp_size=8
Memory usage	OOM on GPU0	Evenly partitioned across 8 GPUs
Streaming behavior	Cutoff after ~30 tokens	Full generation up to max_new_tokens

Verification

After deploying the updated service, run the following sanity check:

curl -X POST http://model-service.local/generate \
  -H "Content-Type: application/json" \
  -d '{"prompt":"Explain the difference between TCP and UDP.", "max_new_tokens":200}'

Expected JSON response (truncated for brevity):

{
  "generated_text": "TCP (Transmission Control Protocol) is a connection‑oriented protocol that provides reliable, ordered delivery of a stream of bytes. UDP (User Datagram Protocol) is connectionless, offering best‑effort delivery without guaranteeing order or reliability..."
}

Additional validation steps:

Monitor GPU memory with nvidia-smi – each GPU should stay below 80 % of VRAM.
Check the streaming endpoint logs for the absence of “token cutoff” warnings.
Run a load test (e.g., hey -n 1000 -c 20) and verify that latency remains stable and no OOM events appear in journalctl -u model-service.

Prevention and Best Practices

Always export an inference‑ready checkpoint after ZeRO‑3 training. Use deepspeed checkpoint --zero3_to_fp16 or the equivalent API call.
Keep training and inference pipelines separate. Do not reuse the raw ZeRO‑3 sharded checkpoint for serving.
Enable DeepSpeed inference (or equivalent tensor parallelism) on the serving side to match the partitioning strategy used during training.
Validate checkpoint metadata before loading. A quick torch.load(.../meta.pt) can confirm the zero_stage value.
Automate health checks that verify the number of generated tokens matches max_new_tokens. Alert on premature termination.
Document the checkpoint conversion step in CI/CD pipelines to avoid accidental deployment of training‑only artifacts.

Related Topic Hub: LLM Systems Troubleshooting Hub

FAQ

Why does the model work locally with device_map='auto' but fail in production?
Because a single‑GPU environment can accommodate the full model, whereas the production node attempts to load the same sharded checkpoint without re‑partitioning, causing an OOM on the first GPU.
Can I skip the deepspeed checkpoint --zero3_to_fp16 step and load the sharded files directly?
Yes, but you must also initialize inference with deepspeed.init_inference and provide the original ZeRO‑3 checkpoint directory. The API will then reconstruct the partitioned model across GPUs.
What if I need to keep optimizer state for downstream fine‑tuning?
Store optimizer checkpoints separately (e.g., optimizer.pt) and only copy the model.pt produced by the conversion step to the inference environment.
Is the token cutoff related to the eos_token_id?
Only indirectly. When the model is partially loaded, the generation loop cannot compute logits for the final layers, causing an early exit before EOS is reached. The observed cutoff length is therefore a symptom, not the root cause.
How do I monitor that the inference service is using all 8 GPUs?
Use deepspeed.utils.logging to emit a log line with the tensor parallel rank and device ID at startup, and verify with nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv that memory usage is distributed.