环境变量
通过环境变量配置 AITraining 行为。
身份验证
Hugging Face
export HF_TOKEN="hf_xxxxxxxxxxxxxxxxxxxxx"
需要用于:
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_LEVEL | 3 | 抑制 TensorFlow 警告 |
TOKENIZERS_PARALLELISM | false | 禁用标记器并行警告 |
BITSANDBYTES_NOWELCOME | 1 | 抑制 bitsandbytes 欢迎消息 |
AUTOTRAIN_TUI_MODE | 1 | 在 TUI 模式下运行时设置 |
调试变量
| 变量 | 默认值 | 用途 |
|---|
PAUSE_ON_FAILURE | 1 | 训练失败时暂停 Space(用于 Spaces 后端) |
优先级
使用 --config 时,配置文件完全控制训练参数。环境变量用于身份验证和系统设置(如 HF_TOKEN、WANDB_API_KEY、CUDA_VISIBLE_DEVICES)。
下一步