# 1. 初始化分布式训练环境

accelerator = Accelerator()
device = accelerator.device

# 在训练脚本中加载DeepSpeed配置
if accelerator.state.deepspeed_plugin is not None:
    ds_config = accelerator.state.deepspeed_plugin.deepspeed_config
    logger.info(f"Loaded DeepSpeed config: {ds_config}")

model_path = get_model_name_path(args.model_name)
if not os.path.exists(model_path):
    raise f"model path: {model_path} not exist."
    logger.info(f"get model path: {model_path}")

# 2. 模型定义
logger.info(f"load model ...")
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map='auto'
    )

model.enable_input_require_grads()  # 关键修复，针对fp16训练，梯度裁剪和梯度检查，防止embedding层也冻结，导致无法训练

peft_config = get_finetune_config(args.model_name, 'lora')
if not peft_config:
    raise "get peft config error."
model = get_peft_model(model, peft_config)
logger.info(f"finetune model:\n {model}")
model.print_trainable_parameters()
model.to(device)

# 加载tokenizer
logger.info(f"load tokenizer ...")
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)

# 3. 数据加载器
logger.info(f"load dataset ...")
dataset = process_data(tokenizer, args.data_path)
dataset.set_format("torch")
train_dataset = dataset.shuffle(seed=42).select(range(1000))
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=args.batch_size)

# 4. 分布式组件准备
model, optimizer, train_loader = accelerator.prepare(
    model,
    torch.optim.AdamW(model.parameters(), lr=args.learning_rate),
    train_dataloader
    )

# 5. 训练控制参数
total_epochs = args.epochs
total_steps = len(train_loader) * total_epochs
start_time = time.time()
current_step = 0

# 6. 创建保存目录（主进程专有）
output_dir = os.path.join(args.output_dir, args.model_name, args.method)
if accelerator.is_main_process:
    save_dir = output_dir+f"_{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    os.makedirs(save_dir, exist_ok=True)

# 7. 分布式训练循环
for epoch in range(total_epochs):
    # 训练阶段
    model.train()
    for batch in train_loader:
      ...