File size: 4,251 Bytes
aa7758f | 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 | #!/usr/bin/env python3
"""Launch the released four-GPU Q-Prefer D2 training recipe."""
from __future__ import annotations
import argparse
import os
import subprocess
import sys
from pathlib import Path
from qprefer_reward.constants import BASE_MODEL_ID, BASE_MODEL_REVISION
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--manifest", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
parser.add_argument("--media-root", type=Path)
parser.add_argument(
"--path-prefix-map",
action="append",
default=[],
metavar="OLD=NEW",
help="Relocate absolute paths from the released manifest",
)
parser.add_argument("--base-model", default=BASE_MODEL_ID)
parser.add_argument("--base-revision", default=BASE_MODEL_REVISION)
parser.add_argument("--initial-special-embeddings", type=Path)
parser.add_argument("--gpus", default="0,1,2,3", help="CUDA indices, comma separated")
parser.add_argument("--master-port", type=int, default=29541)
parser.add_argument("--effective-batch-size", type=int, default=32)
parser.add_argument("--per-device-batch-size", type=int, default=2)
parser.add_argument("--max-steps", type=int, default=1203)
parser.add_argument("--max-train-samples", type=int)
parser.add_argument("--dataloader-num-workers", type=int, default=6)
parser.add_argument("--no-deepspeed", action="store_true")
parser.add_argument("--no-check-media", action="store_true")
parser.add_argument("--resume-from-checkpoint")
return parser.parse_args()
def main() -> None:
args = parse_args()
gpus = [value.strip() for value in args.gpus.split(",") if value.strip()]
if not gpus:
raise ValueError("--gpus must contain at least one CUDA index")
denominator = len(gpus) * args.per_device_batch_size
if args.effective_batch_size % denominator:
raise ValueError(
"effective batch size must be divisible by world size * per-device batch size"
)
accumulation = args.effective_batch_size // denominator
repository = Path(__file__).resolve().parents[1]
command = [
sys.executable,
"-m",
"torch.distributed.run",
f"--nproc_per_node={len(gpus)}",
f"--master_port={args.master_port}",
"-m",
"qprefer_reward.training.train",
"--train-manifest",
str(args.manifest),
"--output-dir",
str(args.output),
"--model-name-or-path",
args.base_model,
"--base-revision",
args.base_revision,
"--max-steps",
str(args.max_steps),
"--per-device-train-batch-size",
str(args.per_device_batch_size),
"--gradient-accumulation-steps",
str(accumulation),
"--dataloader-num-workers",
str(args.dataloader_num_workers),
]
if args.media_root:
command.extend(("--media-root", str(args.media_root)))
for mapping in args.path_prefix_map:
command.extend(("--path-prefix-map", mapping))
if args.initial_special_embeddings:
command.extend(
("--initial-special-embeddings", str(args.initial_special_embeddings))
)
if args.max_train_samples:
command.extend(("--max-train-samples", str(args.max_train_samples)))
if args.no_check_media:
command.append("--no-check-media")
if args.resume_from_checkpoint:
command.extend(("--resume-from-checkpoint", args.resume_from_checkpoint))
if not args.no_deepspeed:
command.extend(
(
"--deepspeed",
str(repository / "training/configs/deepspeed_zero2_no_offload.json"),
)
)
environment = dict(os.environ)
environment["CUDA_VISIBLE_DEVICES"] = ",".join(gpus)
environment.setdefault("FORCE_QWENVL_VIDEO_READER", "decord")
environment.setdefault("VIDEO_IO_MAX_RETRIES", "0")
environment.setdefault("TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC", "1800")
print("Launching:", " ".join(command), flush=True)
subprocess.run(command, cwd=repository, env=environment, check=True)
if __name__ == "__main__":
main()
|