Cccccz/HY

1#!/usr/bin/env python32"""Eight-GPU training entry point for the project-local Predictor-DisCa."""3 4from __future__ import annotations5 6import sys7 8from train_predictor_v4 import main9 10 11def _has_option(name: str) -> bool:12    return any(argument == name or argument.startswith(f"{name}=") for argument in sys.argv[1:])13 14 15if __name__ == "__main__":16    defaults = (17        ("--predictor_variant", "disca"),18        (19            "--output_dir",20            (21                "checkpoints/predictor_disca_blocks0-53_bf16_8gpu_"22                "flr1e-4_blr1e-5_schedsteps2000_swarmup100_bwarmup100_swanlab"23            ),24        ),25        ("--source_blocks", "0,53"),26        ("--swanlab_project", "HY-WorldPlay-Predictor-DisCa"),27        (28            "--swanlab_experiment",29            (30                "predictor-disca-blocks0-53-bf16-8gpu-flr1e-4-blr1e-5-"31                "schedsteps2000-swarmup100-bwarmup100"32            ),33        ),34    )35    injected: list[str] = []36    for option, value in defaults:37        if not _has_option(option):38            injected.extend((option, value))39    if not _has_option("--bucket_by_context") and not _has_option("--no-bucket_by_context"):40        injected.append("--no-bucket_by_context")41    sys.argv[1:1] = injected42    main()43