Heuristic:Hiyouga LLaMA Factory CUDA Memory Optimization

Knowledge Sources	LLaMA-Factory DDP Optimization
Domains	Memory_Optimization, Performance_Optimization
Last Updated	2026-02-06 20:00 GMT

Overview

CUDA memory allocation optimizations automatically applied during distributed training to reduce fragmentation and improve throughput.

Description

LLaMA Factory automatically applies two PyTorch CUDA memory optimizations during distributed training when the OPTIM_TORCH environment variable is enabled (default: 1). These optimizations reduce CUDA memory fragmentation and avoid unnecessary stream synchronization, leading to more efficient memory usage and faster training. The framework also implements explicit garbage collection and cache clearing across all device types (CUDA, NPU, XPU, MPS).

Usage

Use this heuristic automatically during any distributed training. The optimizations are enabled by default. Set OPTIM_TORCH=0 only if you experience compatibility issues.

The Insight (Rule of Thumb)

• Action 1: Keep OPTIM_TORCH=1 (default) for distributed training. This sets:
  • PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True — Uses expandable memory segments to reduce fragmentation.
  • TORCH_NCCL_AVOID_RECORD_STREAMS=1 — Avoids recording NCCL streams to reduce memory overhead.
• Action 2: Call torch_gc() (available via extras.misc) after large operations to free device memory explicitly.
• Trade-off: Expandable segments may slightly increase peak memory in rare cases but significantly reduces OOM errors from fragmentation. Avoiding NCCL record streams reduces memory but may affect debugging.

Reasoning

DDP CUDA optimization from src/llamafactory/launcher.py:80-83:

if is_env_enabled("OPTIM_TORCH", "1"):
    # optimize DDP, see https://zhuanlan.zhihu.com/p/671834539
    env["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
    env["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"

Cross-device memory cleanup from src/llamafactory/extras/misc.py:281-291:

def torch_gc() -> None:
    r"""Collect the device memory."""
    gc.collect()
    if is_torch_xpu_available():
        torch.xpu.empty_cache()
    elif is_torch_npu_available():
        torch.npu.empty_cache()
    elif is_torch_mps_available():
        torch.mps.empty_cache()
    elif is_torch_cuda_available():
        torch.cuda.empty_cache()

The expandable segments feature was introduced in PyTorch to address a common problem: the default CUDA allocator uses fixed-size blocks, leading to fragmentation when training large models. Expandable segments allow the allocator to grow blocks as needed, significantly reducing the chance of OOM errors when total free memory is sufficient but fragmented.

The NCCL stream avoidance optimization reduces memory usage in DDP by not recording the NCCL communication streams in PyTorch's CUDA memory tracking. This is safe because NCCL manages its own memory, and avoiding the recording prevents PyTorch from holding references to already-freed buffers.