Function main

axlearn/cloud/gcp/examples/colocated_python_benchmark.py:230–322

()

def main():
    parser = argparse.ArgumentParser(description="Preload model from GCS checkpoint")
    parser.add_argument(
        "--ckpt_path",
        required=True,
        help="GCS path to checkpoint directory (e.g., gs://bucket/path/to/checkpoint)",
    )
    parser.add_argument(
        "--method",
        choices=["colocated", "default"],
        default="colocated",
        help="Loading method to benchmark: 'colocated' (CPU preload) or 'default' (direct to TPU)",
    )
    parser.add_argument(
        "--profile",
        action="store_true",
        help="Enable JAX profiler (adds overhead, disable for accurate benchmarking)",
    )
    parser.add_argument(
        "--num_iters",
        type=int,
        default=1,
        help="Number of times to repeat the load benchmark (default: 1)",
    )
    args = parser.parse_args()

    # Disable persistent compilation cache for fair benchmarking
    # This ensures benchmarks compile fresh and don't benefit from cached kernels
    os.environ["JAX_ENABLE_COMPILATION_CACHE"] = "0"

    if os.getenv("JAX_PLATFORMS") == "proxy":
        pathwaysutils.initialize()
    else:
        jax.distributed.initialize()

    print(f"JAX devices: {jax.devices()}")

    # Validate checkpoint path
    if not args.ckpt_path.startswith("gs://"):
        raise ValueError(f"Only GCS paths (gs://) are supported, got: {args.ckpt_path}")
    profile_dir = None
    if args.profile:
        # Create timestamped profile directory (minute-level granularity)
        timestamp = datetime.now().strftime("%Y%m%d%H%M")
        base_path = args.ckpt_path.split("/checkpoints")[0]
        profile_dir = f"{base_path}/profiles/{args.method}_{timestamp}/"
        print(f"Profiling enabled - results will be saved to {profile_dir}")

    step = parse_step_from_dir(args.ckpt_path)
    print(f"Starting model preload from: {args.ckpt_path} (step {step})")

    # Read checkpoint structure (doesn't need mesh)
    print("Reading checkpoint structure...")
    state_spec = create_state_spec_from_checkpoint(args.ckpt_path)
    print(f"Found {len(jax.tree_util.tree_leaves(state_spec))} tensors in checkpoint")

    num_iterations = args.num_iters
    print(f"--- Running {args.method} benchmark ({num_iterations} iterations) ---")