Implementation:Axolotl ai cloud Axolotl Prepare Optim Env

Knowledge Sources	Axolotl HuggingFace Accelerate
Domains	Distributed_Training, Infrastructure
Last Updated	2026-02-06 23:00 GMT

Overview

Concrete tool for configuring distributed training environment variables for FSDP and DeepSpeed provided by the Axolotl framework.

Description

The prepare_optim_env function is the top-level entry point for distributed environment setup. It delegates to setup_fsdp_envs for FSDP configuration and setup_deepspeed_env for DeepSpeed configuration. It also handles mixed precision settings, torch.compile configuration, and other optimization-related environment variables.

setup_fsdp_envs sets environment variables like ACCELERATE_USE_FSDP, FSDP_VERSION, FSDP_SHARDING_STRATEGY, etc. based on the YAML config.

setup_deepspeed_env loads the DeepSpeed JSON config, initializes the HfTrainerDeepSpeedConfig, and sets up the distributed state.

Usage

Called early in the training pipeline, after config validation and before model/data loading. Environment variables must be set before Accelerate initializes its distributed state.

Code Reference

Source Location

Repository: axolotl
File: src/axolotl/utils/trainer.py
Lines: L655-688 (prepare_optim_env), L601-630 (setup_fsdp_envs), L545-598 (setup_deepspeed_env)

Signature

def prepare_optim_env(cfg: DictDefault) -> None:
    """Configure distributed training environment.

    Args:
        cfg: Config with fsdp, fsdp_config, deepspeed, fp8, bf16, fp16,
             torch_compile settings.
    """

def setup_fsdp_envs(cfg: DictDefault) -> None:
    """Set FSDP-specific environment variables.

    Args:
        cfg: Config with fsdp_version, fsdp_config (activation_checkpointing,
             offload_params, sync_module_states, auto_wrap_policy, etc.).
    """

def setup_deepspeed_env(
    cfg: DictDefault,
    stage: int | None = None,
) -> None:
    """Set DeepSpeed environment and initialize config.

    Args:
        cfg: Config with deepspeed (file path or inline dict).
        stage: Optional ZeRO stage override (1/2/3).
    """

Import

from axolotl.utils.trainer import prepare_optim_env, setup_fsdp_envs, setup_deepspeed_env

I/O Contract

Inputs

Name	Type	Required	Description
cfg	DictDefault	Yes	Config with distributed training settings (fsdp_config, deepspeed, etc.)

Outputs

Name	Type	Description
(side effects)	Environment variables	Sets ACCELERATE_USE_FSDP, FSDP_VERSION, FSDP_SHARDING_STRATEGY, and DeepSpeed-related env vars

Usage Examples

FSDP Setup

from axolotl.utils.trainer import prepare_optim_env

# cfg.fsdp_config = {
#     "fsdp_version": 2,
#     "activation_checkpointing": True,
#     "auto_wrap_policy": "TRANSFORMER_BASED_WRAP",
#     "transformer_layer_cls_to_wrap": "LlamaDecoderLayer",
# }
prepare_optim_env(cfg)
# Environment variables now set for FSDP2 with activation checkpointing

DeepSpeed Setup

from axolotl.utils.trainer import prepare_optim_env

# cfg.deepspeed = "deepspeed_configs/zero3_bf16.json"
prepare_optim_env(cfg)
# DeepSpeed ZeRO-3 environment configured

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment