| |
| """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() |
|
|
