# Best Practices for NVFP4 Checkpoint Loading and Weight Materialization in LongLive

> Optimize NVFP4 checkpoint loading and weight materialization in LongLive. Safely load checkpoints, unwrap generators, clean FSDP prefixes, and drop master weights to save GPU memory.

- Repository: [NVIDIA Research Projects/LongLive](https://github.com/NVlabs/LongLive)
- Tags: best-practices
- Published: 2026-05-24

---

**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`](https://github.com/NVlabs/LongLive/blob/main/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:

```python
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:

```python

# 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`](https://github.com/NVlabs/LongLive/blob/main/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

```python
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:

```python
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`](https://github.com/NVlabs/LongLive/blob/main/utils/nvfp4_checkpoint.py)**: Core loading utilities including format detection, key cleaning, and master weight dropping
- **[`utils/inference_utils.py`](https://github.com/NVlabs/LongLive/blob/main/utils/inference_utils.py)**: Higher-level inference wrappers that internally orchestrate these helpers
- **[`scripts/save_merged_nvfp4_generator.py`](https://github.com/NVlabs/LongLive/blob/main/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`](https://github.com/NVlabs/LongLive/blob/main/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`](https://github.com/NVlabs/LongLive/blob/main/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.