Best Practices for NVFP4 Checkpoint Loading and Weight Materialization in LongLive

Load NVFP4 checkpoints safely by detecting format types with is_nvfp4_state_dict, unwrapping generator state via unwrap_generator_state_dict, cleaning FSDP prefixes with clean_fsdp_state_dict_keys, and dropping high-precision master weights using drop_fouroversix_master_weights to minimize GPU memory while preserving materialized low-precision buffers.

LongLive stores generator checkpoints in a quantized NVFP4 format that contains pre-materialized low-precision weight buffers. Proper NVFP4 checkpoint loading and weight materialization in LongLive requires specific utilities from utils/nvfp4_checkpoint.py to handle format detection, EMA wrappers, and memory optimization. Following these canonical steps ensures efficient inference and fine-tuning without GPU OOM errors or key mismatches.

Understanding NVFP4 Checkpoint Architecture

NVFP4 checkpoints in LongLive store already-quantized weight buffers rather than high-precision master copies. This design reduces storage footprint and accelerates inference, but requires careful handling during load operations to ensure compatibility with different model wrappers and distributed training setups.

Detecting Checkpoint Formats

Before loading, verify the checkpoint contains materialized NVFP4 buffers using format detection utilities. The is_nvfp4_state_dict(state_dict) function checks for the presence of quantized_weight_values keys, while is_te_nvfp4_checkpoint(chkpt) detects Transformer-Engine wrapped variants.

According to the LongLive source code, these helpers prevent accidental loading of incompatible formats:

• is_nvfp4_state_dict (lines 23-28): Validates materialized NVFP4 buffers exist
• is_te_nvfp4_checkpoint (lines 30-35): Detects TE-specific checkpoint layouts

Always validate the checkpoint type before proceeding with generator extraction.

Handling EMA and FSDP Wrappers

LongLive checkpoints often wrap the generator inside multiple possible keys: generator, generator_ema, or model. The unwrap_generator_state_dict(chkpt, use_ema) utility (lines 38-48) normalizes these layouts, returning a flat state dictionary ready for loading.

For checkpoints saved with Fully-Sharded Data-Parallel (FSDP), keys may contain the prefix "_fsdp_wrapped_module.". The clean_fsdp_state_dict_keys(state_dict) helper (lines 50-53) strips these prefixes automatically, preventing RuntimeError exceptions during load_state_dict operations.

Step-by-Step Loading Pipeline

Execute these steps sequentially to safely load NVFP4 checkpoints while managing memory efficiently.

Load to CPU and Extract Generator State

Always load checkpoints to CPU first to prevent GPU OOM during the initial load phase. Use map_location='cpu' in torch.load, then extract the generator state:

import torch
from utils.nvfp4_checkpoint import (
    is_nvfp4_state_dict,
    unwrap_generator_state_dict,
    clean_fsdp_state_dict_keys,
)

# Load to CPU to prevent GPU memory overflow

raw_ckpt = torch.load(ckpt_path, map_location="cpu")

# Normalize wrapper layouts (handles EMA vs non-EMA)

gen_state = unwrap_generator_state_dict(raw_ckpt, use_ema=config.use_ema)

# Verify NVFP4 format

assert is_nvfp4_state_dict(gen_state), "Checkpoint missing NVFP4 buffers"

Clean FSDP Prefixes

If the checkpoint originated from FSDP training, clean the keys before loading:

# Remove "_fsdp_wrapped_module." prefixes

gen_state = clean_fsdp_state_dict_keys(gen_state)

This step is idempotent and safe to run on non-FSDP checkpoints, making it a robust addition to any loading pipeline.

Materialize for Inference vs Fine-tuning

The materialization strategy depends on your use case:

• For inference: Quantize the model with materialize_for_inference=False, load the state dict, then drop master weights
• For fine-tuning: Use quantize_model_for_fouroversix_nvfp4 with materialize_for_inference=False before loading to enable gradient updates on low-precision buffers

The quantization wrapper replaces eligible modules with NVFP4-compatible layers that understand the quantized_weight_values buffer format.

Memory Optimization Strategies

NVFP4 checkpoints store materialized low-precision tensors, rendering original high-precision master weights redundant for inference.

Dropping Master Weights

After model.load_state_dict(state), immediately invoke drop_fouroversix_master_weights(model) (lines 34-64 in utils/nvfp4_checkpoint.py). This utility:

• Removes high-precision master weight copies
• Clears internal quantization caches
• Returns a list of module names affected for verification

from utils.nvfp4_checkpoint import drop_fouroversix_master_weights

# Load materialized buffers

model.load_state_dict(gen_state, strict=False)

# Free GPU memory by removing redundant master weights

dropped_modules = drop_fouroversix_master_weights(model)
print(f"Optimized memory by dropping master weights from: {dropped_modules}")

This step is critical for deployment scenarios where GPU memory is constrained.

Complete Implementation Example

Combine all utilities into a robust loading function for inference:

from utils.nvfp4_checkpoint import (
    is_nvfp4_state_dict,
    unwrap_generator_state_dict,
    clean_fsdp_state_dict_keys,
    drop_fouroversix_master_weights,
)
from utils.quant import ModelQuantizationConfig, quantize_model_with_filter

def load_nvfp4_generator(ckpt_path: str, config, device: str = "cuda"):
    # 1. Load raw checkpoint to CPU

    raw_ckpt = torch.load(ckpt_path, map_location="cpu")
    
    # 2. Extract generator (handles EMA/model wrappers)

    gen_state = unwrap_generator_state_dict(raw_ckpt, use_ema=config.use_ema)
    
    # 3. Verify NVFP4 format

    if not is_nvfp4_state_dict(gen_state):
        raise ValueError("Checkpoint does not contain NVFP4 buffers")
    
    # 4. Clean FSDP prefixes

    gen_state = clean_fsdp_state_dict_keys(gen_state)
    
    # 5. Configure quantization

    quant_cfg = ModelQuantizationConfig(
        scale_rule=getattr(config, "model_quant_scale_rule", "static_6"),
        quantize_backend=getattr(config, "model_quant_backend", None),
    )
    
    # 6. Wrap model with NVFP4 layers (non-materialized)

    model = quantize_model_with_filter(
        model,
        quant_config=quant_cfg,
        filtered_modules=getattr(config, "model_quant_filtered_modules", None),
        use_default_filtered_modules=getattr(
            config, "model_quant_use_default_filtered_modules", True
        ),
        materialize_for_inference=False,
    )
    
    # 7. Load state dict and optimize memory

    model.load_state_dict(gen_state, strict=False)
    drop_fouroversix_master_weights(model)
    
    # 8. Move to device and set eval mode

    model.to(device)
    model.eval()
    return model

Key Files and Utilities

The following files contain the canonical implementations referenced above:

• utils/nvfp4_checkpoint.py: Core loading utilities including format detection, key cleaning, and master weight dropping
• utils/inference_utils.py: Higher-level inference wrappers that internally orchestrate these helpers
• scripts/save_merged_nvfp4_generator.py: Reference implementation for saving merged NVFP4 generators, useful for understanding the checkpoint format structure

Summary

• Detect format using is_nvfp4_state_dict before processing to ensure compatibility with materialized low-precision buffers
• Normalize wrappers via unwrap_generator_state_dict to handle EMA, generator, and model key variations consistently
• Clean FSDP keys with clean_fsdp_state_dict_keys to remove distributed training prefixes that break state dict loading
• Load to CPU first using map_location='cpu' to prevent GPU memory overflow during checkpoint initialization
• Drop master weights immediately after loading via drop_fouroversix_master_weights to minimize GPU memory footprint for inference
• Materialize appropriately by setting materialize_for_inference=False during quantization setup, enabling proper buffer interpretation without unnecessary memory allocation

Frequently Asked Questions

How do I detect if a LongLive checkpoint uses NVFP4 format?

Use the is_nvfp4_state_dict(state_dict) utility from utils/nvfp4_checkpoint.py. This function checks for the presence of quantized_weight_values keys within the state dictionary, confirming the checkpoint contains pre-materialized NVFP4 buffers rather than high-precision weights. For Transformer-Engine variants, use is_te_nvfp4_checkpoint(chkpt) instead.

What is the purpose of unwrap_generator_state_dict?

This helper normalizes different checkpoint wrapper layouts by extracting the generator state from various possible container keys like generator, generator_ema, or model. It accepts a use_ema boolean parameter to select between EMA and standard generator weights, ensuring consistent state dictionary formatting regardless of how the checkpoint was saved during training.

When should I drop high-precision master weights?

Drop master weights immediately after calling model.load_state_dict() and before moving the model to the GPU for inference. Invoke drop_fouroversix_master_weights(model) to remove redundant high-precision copies, as NVFP4 checkpoints already contain materialized low-precision buffers. This step is essential for minimizing GPU memory usage during deployment.

How do I handle FSDP-prefixed keys in NVFP4 checkpoints?

Run clean_fsdp_state_dict_keys(state_dict) from utils/nvfp4_checkpoint.py on your state dictionary before loading. This utility strips the "_fsdp_wrapped_module." prefix that Fully-Sharded Data-Parallel training adds to parameter keys, preventing key mismatch errors during load_state_dict operations on non-FSDP or differently-sharded model instances.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}