noao-vlm-0 train.py 详细分析

torch.manual_seed(0)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(0)

os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
PIL.PngImagePlugin.MAX_TEXT_CHUNK = 100 * 1024 * 1024

def dist_gather(obj):
    result = [None] * dist.get_world_size()
    dist.all_gather_object(result, obj, group=PG_CPU)  # 使用 CPU 组
    return result

def get_lr(it, max_lr, max_steps):
    min_lr = max_lr * 0.1
    warmup_steps = max_steps * 0.03
    
    # Phase 1: Linear Warmup
    if it < warmup_steps:
        return max_lr * (it+1) / warmup_steps
    
    # Phase 2: Cosine Decay
    if it <= max_steps:
        decay_ratio = (it - warmup_steps) / (max_steps - warmup_steps)
        coeff = 0.5 * (1.0 + math.cos(math.pi * decay_ratio))
        return min_lr + coeff * (max_lr - min_lr)
    
    # Phase 3: Minimum LR
    return min_lr

if train_cfg.resume_from_vlm_checkpoint:
    model = VisionLanguageModel.from_pretrained(vlm_cfg.vlm_checkpoint_path)
else:
    model = VisionLanguageModel(vlm_cfg, load_backbone=True)

param_groups = []

# MP 层：高学习率（新初始化）
if train_cfg.lr_mp > 0:
    param_groups.append({'params': model.MP.parameters(), 'lr': 0.00512})
else:
    freeze_parameters(model.MP)

# Vision 骨干：低学习率（预训练）
if train_cfg.lr_vision_backbone > 0:
    param_groups.append({'params': model.vision_encoder.parameters(), 'lr': 5e-5})
else:
    freeze_parameters(model.vision_encoder)

# Language 骨干：低学习率（预训练）
if train_cfg.lr_language_backbone > 0:
    param_groups.append({'params': model.decoder.parameters(), 'lr': 5e-5})
else:
    freeze_parameters(model.decoder)

optimizer = optim.AdamW(param_groups)

# 编译优化（可选）
if train_cfg.compile:
    model = torch.compile(model)

# 分布式包装
if is_dist():
    model = wrap_model(model)  # DistributedDataParallel

while global_step < max_training_steps:
    epoch += 1
    for i, batch in enumerate(synchronized_dataloader_step(iter_train_loader)):
        # 梯度累积逻辑
        is_update_step = (i + 1) % gradient_accumulation_steps == 0
        
        # 前向传播
        # 反向传播
        # 优化器更新（仅 is_update_step=True）

if is_dist() and gradient_accumulation_steps > 1 and not is_update_step:
    context = model.no_sync()  # 跳过同步
else:
    context = contextlib.nullcontext()  # 正常同步

accumulated_stats = {
    'tokens_per_second': [],      # 吞吐量
    'data_load_time': [],          # 数据加载时间
    'fw_bw_time': [],              # 前向+反向时间
    'post_process_time': [],       # 后处理时间
    'images_per_sample': [],       # 每样本图像数
}

if eval_in_epochs and global_step % eval_interval == 0 and is_update_step:

checkpoint_path = f"{vlm_checkpoint_path}/{run_name}/step_{global_step}"
save_model.save_pretrained(checkpoint_path)

if avg_val_loss < best_val_loss:
    best_val_loss = avg_val_loss
    best_model_path = checkpoint_path_step

# 梯度累积时跳过中间步骤的同步
if is_dist() and gradient_accumulation_steps > 1 and not is_update_step:
    context = model.no_sync()
else:
    context = contextlib.nullcontext()

# NCCL 进程组：GPU 计算和梯度同步
dist.init_process_group(backend='nccl')

# Gloo 进程组：CPU 数据收集
PG_CPU = dist.new_group(backend="gloo")

def synchronized_dataloader_step(iterator, is_distributed):
    """确保所有 rank 在每步同步加载数据"""
    if is_distributed:
        # 同步机制防止某些 rank 提前结束
        pass
    return iterator

# 关键同步点
if is_dist():
    dist.barrier(device_ids=[local_rank])

autocast_context = torch.autocast(
    device_type=device.type,
    dtype=torch.bfloat16 if device.type in ['cuda', 'cpu'] else torch.float16
)

gradient_accumulation_steps = 8
effective_batch_size = batch_size * world_size * gradient_accumulation_steps
# 2 × 8 × 8 = 128

# 配置
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'

# 验证前清空缓存
if device == "cuda":
    torch.cuda.empty_cache()

if train_cfg.compile:
    model = torch.compile(model)

train_loader = DataLoader(
    dataset,
    batch_size=2,
    num_workers=3,           # 多进程加载
    pin_memory=True,         # 固定内存
    persistent_workers=False, # 节省内存
    drop_last=True,          # 保证批次一致
)

train_ds = load_dataset(path, streaming=True)

run = wandb.init(
    entity="HuggingFace",
    project="nanoVLM",
    config={
        "VLMConfig": asdict(vlm_cfg),
        "TrainConfig": asdict(train_cfg)
    },
    name=run_name,
)

lmms_eval_step = "<lmms-eval-step>"
run.define_metric(name="lmms_eval/*", step_metric=lmms_eval_step)

if use_lmms_eval and global_step % (eval_interval*2) == 0:
    cmd = f"sbatch eval.slurm {checkpoint_path} {global_step} {run_name} ..."
    subprocess.run(cmd, shell=True)

lmms_eval_tasks = 'mmstar,mmmu_val,ocrbench,textvqa_val,docvqa_val,
                   scienceqa,mme,infovqa_val,chartqa'

# 监视评估结果目录
eval_results_dir = os.path.join('eval_results', run_name)

# 解析新结果
for result_file in os.listdir(eval_results_dir):
    match = re.fullmatch(r"step_(\d+)\.json", result_file)
    if match and step not in logged_eval_steps:
        # 加载并记录到 wandb
        run.log(metrics, step=global_step)
        logged_eval_steps.add(step)

run_name = (f"nanoVLM_{vit}_{mp}_{llm}_{num_gpus}_"
            f"{batch_size}_{max_steps}_{lr_info}_{date}")

# 所有 rank 收集数据
accumulated_stats['tokens_per_second'].append(local_value)

# 所有 rank 参与聚合
if is_dist():
    all_values = dist_gather(accumulated_stats['tokens_per_second'])
    all_values_flat = [item for sublist in all_values for item in sublist]
    global_avg = mean(all_values_flat)
else:
    global_avg = mean(accumulated_stats['tokens_per_second'])

# 仅 master 记录
if is_master():
    run.log({"avg_tokens_per_second": global_avg})

for dataset_name in dataset_names_to_load:
    try:
        train_ds = load_dataset(...)
        combined_train_data.append(train_ds)
    except Exception as e:
        print(f"Warning: Failed to load {dataset_name}. Error: {e}")
        continue  # 跳过失败的数据集

if not combined_train_data:
    raise ValueError("No valid datasets were loaded.")

try:
    with open(result_file, 'r') as f:
        eval_data = json.load(f)
    # 处理数据
except (ValueError, KeyError, json.JSONDecodeError) as e:
    print(f"Warning: Could not process {result_file}. Error: {e}")
    continue

# 降低批次大小
train_cfg.batch_size = 1

# 降低图像数量
train_cfg.max_images_per_knapsack = 10

# 降低序列长度
vlm_cfg.lm_max_length = 2048

# 降低学习率
train_cfg.lr_mp = 0.001

# 增加 warmup
# 在 get_lr() 中调整 warmup_steps 比例

# 增加 worker 数量
num_workers = 6

# 启用流式数据集
train_cfg.stream_dataset = True

# 启用编译优化
train_cfg.compile = True

# 使用更大的有效批次
train_cfg.gradient_accumulation_steps = 16

# 优化 DataLoader
num_workers = 4  # 根据 CPU 核心数调整

# 训练完成后自动推送
if vlm_cfg.hf_repo_name:
    best_model.push_to_hub(vlm_cfg.hf_repo_name)

# 修改 config.py
lmms_eval_tasks = 'mmstar,mmmu_val,your_custom_task'

def custom_lr_schedule(step, max_lr, max_steps):
    # 实现自定义逻辑
    if step < warmup:
        return ...
    elif step < decay_point:
        return ...
    else:
        return ...

# 在训练循环中使用
adj_lr = custom_lr_schedule(global_step, max_lr, max_steps)

# 修改 data/custom_transforms.py
class MyCustomTransform:
    def __call__(self, image):
        # 自定义增强逻辑
        return transformed_image

# 在 get_image_processor() 中添加