File size: 7,135 Bytes
df529cc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | #!/usr/bin/env python3
"""Diagnostic for ZeRO-2 with unequal real ChartQA image-token lengths.
This is a bounded regression diagnostic for the OPD-only training failure where
one rank reached a Qwen MLP gradient all-reduce while a longer-image rank had
not yet reached the same autograd hook. It loads the 0.5B student only, picks
the highest-resolution hard-correct ChartQA rows, and performs a few real VLM
backward/optimizer steps with the bounded production ZeRO-2 JSON. Formal
OPD-only training uses native DDP; this script is retained only to inspect
DeepSpeed behavior when explicitly requested.
"""
from __future__ import annotations
import argparse
import copy
import json
import os
import sys
import time
from pathlib import Path
import deepspeed
import torch
import torch.distributed as dist
from PIL import Image
from transformers import AutoProcessor, LlavaOnevisionForConditionalGeneration
from transformers.models.llava_onevision.modeling_llava_onevision import image_size_to_num_patches
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from config.loader import load_config
from data_utils.chart.data_collector import prepare_chart_rl_data
from data_utils.chart.teacher_gate import filter_chartqa_rows_with_teacher_gate
from data_utils.commom_util import collate_fn
from opsd_utils.teacher_batching import student_batch_num_images_tensor
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--config", default="config/config_opd_only_7b_chartqa.yaml")
parser.add_argument("--iterations", type=int, default=3)
parser.add_argument("--samples-per-rank", type=int, default=2)
return parser.parse_args()
def _deepspeed_config(*, world_size: int, micro_batch_size: int) -> dict:
root = Path(__file__).resolve().parents[1]
payload = json.loads((root / "configs/deepspeed/zero2_bf16.json").read_text(encoding="utf-8"))
payload = copy.deepcopy(payload)
payload.pop("scheduler", None)
payload["optimizer"]["params"] = {
"lr": 1.0e-6,
"betas": [0.9, 0.999],
"eps": 1.0e-8,
"weight_decay": 0.0,
"torch_adam": True,
"adam_w_mode": True,
}
payload["gradient_accumulation_steps"] = 1
payload["gradient_clipping"] = 0.0
payload["train_micro_batch_size_per_gpu"] = micro_batch_size
payload["train_batch_size"] = world_size * micro_batch_size
return payload
def _highest_resolution_rows(config: dict, count: int) -> tuple[list[dict], list[int]]:
rows = prepare_chart_rl_data(config["dataset"]["train_dataset"])
for index, row in enumerate(rows):
row["_dyme_index"] = index
rows, _ = filter_chartqa_rows_with_teacher_gate(rows, config["dataset"]["teacher_gate"])
processor = AutoProcessor.from_pretrained(config["model"]["pretrained_model_path"], local_files_only=True)
grid = processor.image_processor.image_grid_pinpoints
patch_size = processor.image_processor.size["height"]
ranked: list[tuple[int, int, dict]] = []
for index, row in enumerate(rows):
with Image.open(row["image"]) as image:
num_patches = image_size_to_num_patches(image.size, grid, patch_size)
ranked.append((int(num_patches), index, row))
ranked.sort(key=lambda item: (-item[0], item[1]))
chosen = ranked[:count]
return [row for _patches, _index, row in chosen], [patches for patches, _index, _row in chosen]
def main() -> None:
args = parse_args()
if args.iterations <= 0 or args.samples_per_rank <= 0:
raise ValueError("--iterations and --samples-per-rank must be positive")
local_rank = int(os.environ["LOCAL_RANK"])
torch.cuda.set_device(local_rank)
deepspeed.init_distributed(dist_backend="nccl")
rank = dist.get_rank()
world_size = dist.get_world_size()
device = torch.device("cuda", local_rank)
try:
config = load_config(args.config)
candidate_count = world_size * args.samples_per_rank
rows, patch_counts = _highest_resolution_rows(config, candidate_count)
offset = rank * args.samples_per_rank
local_rows = rows[offset : offset + args.samples_per_rank]
local_patch_counts = patch_counts[offset : offset + args.samples_per_rank]
if len(local_rows) != args.samples_per_rank:
raise RuntimeError("not enough ChartQA rows for distributed high-resolution diagnostic")
model_path = config["model"]["pretrained_model_path"]
processor = AutoProcessor.from_pretrained(model_path, local_files_only=True)
processor.tokenizer.padding_side = "left"
model = LlavaOnevisionForConditionalGeneration.from_pretrained(
model_path,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2",
local_files_only=True,
low_cpu_mem_usage=True,
)
model.base_model.vision_tower.requires_grad_(False)
model.config.use_cache = False
model.to(device)
trainable_parameters = sum(parameter.numel() for parameter in model.parameters() if parameter.requires_grad)
engine, _, _, _ = deepspeed.initialize(
model=model,
model_parameters=[parameter for parameter in model.parameters() if parameter.requires_grad],
config=_deepspeed_config(world_size=world_size, micro_batch_size=args.samples_per_rank),
)
bucket_size = int(engine.optimizer.reduce_bucket_size)
batch = collate_fn(local_rows, processor, label_id=None)
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
pixel_values = batch["pixel_values"].to(device=device, dtype=torch.bfloat16)
image_sizes = batch["image_sizes"].to(device)
batch_num_images = student_batch_num_images_tensor(pixel_values, input_ids.size(0))
dist.barrier()
started = time.perf_counter()
for _ in range(args.iterations):
outputs = engine(
input_ids=input_ids,
attention_mask=attention_mask,
pixel_values=pixel_values,
image_sizes=image_sizes,
batch_num_images=batch_num_images,
logits_to_keep=2,
)
loss = outputs.logits.float().mean()
engine.backward(loss)
engine.step()
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - started
print(
"ZeRO-2 VLM bounded-bucket diagnostic passed: "
f"rank={rank}/{world_size} patch_counts={local_patch_counts} "
f"input_shape={tuple(input_ids.shape)} bucket={bucket_size} "
f"trainable={trainable_parameters} iterations={args.iterations} "
f"elapsed_s={elapsed:.3f} loss={float(loss.detach().item()):.6f}",
flush=True,
)
dist.barrier()
finally:
if dist.is_initialized():
dist.destroy_process_group()
if __name__ == "__main__":
main()
|