Implementation:NVIDIA DALI FP16 Utilities
| Knowledge Sources | |
|---|---|
| Domains | Video_Processing, Mixed_Precision |
| Last Updated | 2026-02-08 16:00 GMT |
Overview
Provides FP16 (half-precision) training utilities for PyTorch, including an FP16 model wrapper, an FP16 optimizer with master weight copy in FP32, and dynamic/static loss scaling support.
Description
This module implements the infrastructure needed for mixed-precision (FP16) training in PyTorch for the video super-resolution pipeline. It provides three main components: precision conversion helpers, an FP16 model wrapper, and an FP16 optimizer.
The conversion helpers (`fp32_to_fp16`, `fp16_to_fp32`) recursively traverse nested tuples and lists of tensors, converting between float32 and float16 precision. They correctly handle both raw tensors and `Parameter`/`Variable` types by checking the underlying data type before conversion.
`FP16_Module` is an `nn.Module` wrapper that converts a model's parameters to FP16, automatically handles input conversion from FP32 to FP16 before the forward pass, and converts outputs back to FP32. This allows the rest of the training code to work with FP32 while the model operates in FP16 for GPU memory efficiency and throughput improvements.
`FP16_Optimizer` wraps an existing PyTorch optimizer to enable FP16 training with a master copy of FP32 weights. During initialization, it clones and flattens all FP16 model parameters into contiguous FP32 tensors, replaces the optimizer's parameter groups with these FP32 copies, and maintains references to the original FP16 parameters. The training cycle works as follows: gradients are computed in FP16, copied and flattened into the FP32 parameter space, downscaled by the loss scale factor, the FP32 optimizer step is applied, and the updated FP32 parameters are copied back to the FP16 model parameters. The optimizer supports both static loss scaling (fixed multiplier) and dynamic loss scaling (via `DynamicLossScaler`) which automatically adjusts the scale factor based on gradient overflow detection. It provides gradient clipping, state dict save/load, and closure-based step support.
Usage
Use these utilities in the video super-resolution training pipeline to enable mixed-precision training for improved GPU performance. Wrap the model with `FP16_Module` and the optimizer with `FP16_Optimizer` to transparently handle precision management.
Code Reference
Source Location
- Repository: NVIDIA_DALI
- File: docs/examples/use_cases/video_superres/common/fp16.py
- Lines: 1-424
Signature
def conversion_helper(val, conversion):
"""Apply conversion to val. Recursively apply conversion if val is a nested tuple/list."""
...
def fp32_to_fp16(val):
"""Convert fp32 val to fp16."""
...
def fp16_to_fp32(val):
"""Convert fp16 val to fp32."""
...
class FP16_Module(nn.Module):
def __init__(self, module): ...
def forward(self, *inputs, **kwargs): ...
class FP16_Optimizer(object):
def __init__(self, optimizer, static_loss_scale=1.0, dynamic_loss_scale=False): ...
def zero_grad(self): ...
def step(self, closure=None): ...
def backward(self, loss): ...
def state_dict(self) -> dict: ...
def load_state_dict(self, state_dict): ...
def clip_fp32_grads(self, max_norm, norm_type=2): ...
@property
def loss_scale(self) -> float: ...
Import
from common.fp16 import FP16_Module, FP16_Optimizer
model = FP16_Module(model)
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| module | nn.Module | Yes | PyTorch model to wrap for FP16 inference (FP16_Module) |
| optimizer | torch.optim.Optimizer | Yes | PyTorch optimizer with FP16 parameters (FP16_Optimizer) |
| static_loss_scale | float | No | Fixed loss scale multiplier (default: 1.0) |
| dynamic_loss_scale | bool | No | Enable dynamic loss scaling (default: False; overrides static_loss_scale) |
| loss | torch.Tensor | Yes | Scalar loss tensor for backward() call |
| max_norm | float | No | Maximum gradient norm for clip_fp32_grads() |
Outputs
| Name | Type | Description |
|---|---|---|
| FP16_Module output | torch.Tensor (fp32) | Model output converted back to FP32 |
| state_dict | dict | Optimizer state including loss_scaler, overflow flag, and inner optimizer state |
| grad_norm | float | Total gradient norm after clipping (-1 if overflow occurred) |
Usage Examples
Mixed Precision Training Loop
import torch
from common.fp16 import FP16_Module, FP16_Optimizer
# Wrap model and optimizer
model = FP16_Module(model).cuda()
optimizer = FP16_Optimizer(
torch.optim.Adam(model.parameters(), lr=1e-4),
dynamic_loss_scale=True,
)
# Training loop
for images, targets in dataloader:
optimizer.zero_grad()
# Forward pass (inputs auto-converted to FP16)
outputs = model(images.cuda())
# Compute loss in FP32
loss = criterion(outputs, targets.cuda())
# Backward with loss scaling
optimizer.backward(loss)
# Optional gradient clipping
optimizer.clip_fp32_grads(max_norm=1.0)
# Step (handles overflow detection, FP32->FP16 copy)
optimizer.step()