跳转到主要内容

环境变量

通过环境变量配置 AITraining 行为。

身份验证

Hugging Face

export HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxx"
需要用于:
  • 私有模型
  • 私有数据集
  • 推送到 Hub

Weights & Biases

export WANDB_API_KEY="xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
使用 --log wandb 时需要。

GPU 配置

选择 GPU

# 仅使用 GPU 0
export CUDA_VISIBLE_DEVICES=0

# 使用 GPU 0 和 1
export CUDA_VISIBLE_DEVICES=0,1

# 禁用 GPU(仅 CPU)
export CUDA_VISIBLE_DEVICES=""

内存管理

# 限制内存碎片并启用可扩展段
export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512,expandable_segments:True

GPU 覆盖

# 强制 CPU 训练(覆盖 GPU 检测)
export AUTOTRAIN_FORCE_NUM_GPUS=0

# 强制单 GPU
export AUTOTRAIN_FORCE_NUM_GPUS=1

Mac(Apple Silicon)设置

MPS 控制

# 禁用 MPS(在 Mac 上强制 CPU)
export AUTOTRAIN_DISABLE_MPS=1

# 即使设置不兼容也强制启用 MPS
export AUTOTRAIN_ENABLE_MPS=1
当 MPS 被禁用或不兼容时,AITraining 会自动设置 PYTORCH_ENABLE_MPS_FALLBACK=1 以回退到 CPU。

安全设置

远程代码执行

# 禁用所有模型加载的 trust_remote_code(更安全)
export ALLOW_REMOTE_CODE=false
设置 ALLOW_REMOTE_CODE=false 可能会阻止加载需要自定义代码的模型(如某些多模态模型)。默认值为 true

日志记录

W&B 设置

export WANDB_PROJECT=my-project
export WANDB_ENTITY=my-team
export WANDB_MODE=offline  # 不同步到云端

Hub 设置

默认用户名

export HF_USERNAME=my-username

缓存目录

export HF_HOME=/path/to/cache
export TRANSFORMERS_CACHE=/path/to/cache

训练设置

禁用遥测

export HF_HUB_DISABLE_TELEMETRY=1

离线模式

export TRANSFORMERS_OFFLINE=1
export HF_HUB_OFFLINE=1

使用示例

完整设置脚本

创建设置脚本 setup_env.sh
#!/bin/bash

# 身份验证
export HF_TOKEN="hf_your_token_here"
export WANDB_API_KEY="your_wandb_key"

# GPU 设置
export CUDA_VISIBLE_DEVICES=0,1

# 缓存
export HF_HOME=~/.cache/huggingface

# W&B 项目设置
export WANDB_PROJECT=my-project
在训练之前加载:
source setup_env.sh
aitraining llm --train ...

使用 .env 文件

创建 .env 文件:
HF_TOKEN=hf_xxxxx
WANDB_API_KEY=xxxxx
CUDA_VISIBLE_DEVICES=0
加载方式:
export $(cat .env | xargs)
aitraining --config training.yaml

内部变量

这些由 AITraining 自动设置:
变量用途
TF_CPP_MIN_LOG_LEVEL3抑制 TensorFlow 警告
TOKENIZERS_PARALLELISMfalse禁用标记器并行警告
BITSANDBYTES_NOWELCOME1抑制 bitsandbytes 欢迎消息
AUTOTRAIN_TUI_MODE1在 TUI 模式下运行时设置

调试变量

变量默认值用途
PAUSE_ON_FAILURE1训练失败时暂停 Space(用于 Spaces 后端)

优先级

使用 --config 时,配置文件完全控制训练参数。环境变量用于身份验证和系统设置(如 HF_TOKENWANDB_API_KEYCUDA_VISIBLE_DEVICES)。

下一步