Function main

train/train.py:89–199

(data_args=None, training_args=None)

def main(data_args=None, training_args=None):

    # Initialize Overwatch =>> Wraps `logging.Logger`
    overwatch = initialize_overwatch(__name__)
    
    torch.cuda.set_device(device_id := overwatch.local_rank())
    torch.cuda.empty_cache()
    run_id = (
        f"n{training_args.expected_world_size // 8}+b{training_args.per_device_batch_size}+x{training_args.seed}"
    )


    worker_init_fn = None

    os.makedirs(run_dir := (data_args.run_root_dir / run_id), exist_ok=True)
    os.makedirs(data_args.run_root_dir / run_id / "checkpoints", exist_ok=True)

    model = load_vla(training_args.pretrained_checkpoint, hf_token=training_args.hf_token, load_for_training=True, grid_size=training_args.grid_size)

    for param in model.parameters():
        assert param.dtype == torch.float32, f"Loaded VLM parameter not in full precision: {param}"

    # Determine training "stage" based on frozen vs unfrozen parameters --> supports different fine-tuning schemes!
    if not training_args.freeze_vision_backbone and not training_args.freeze_llm_backbone:
        stage = "vla-full-train"  # Full fine-tuning
    elif training_args.freeze_vision_backbone and not training_args.freeze_llm_backbone:
        stage = "vla-train"  # Frozen vision encoder
    elif not training_args.freeze_vision_backbone and training_args.freeze_llm_backbone:
        assert training_args.unfreeze_last_llm_layer, "You should unfreeze at least the last layer of your LLM!"
        stage = "vla-sandwich-train"  # Fine-tuning vision encoder, projector, and LLM last layer
    elif training_args.freeze_vision_backbone and training_args.freeze_llm_backbone:
        assert training_args.unfreeze_last_llm_layer, "Need to unfreeze at least last LLM layer to train!"
        stage = "vla-last-layer-train"  # Fine-tuning LLM last layer only
    else:
        raise ValueError(
            "Weight freezing configuration not supported. VLA config has the following parameters: "
            f"freeze_vision_backbone: {training_args.freeze_vision_backbone}"
            f"freeze_llm_backbone: {training_args.freeze_llm_backbone}"
            f"unfreeze_last_llm_layer: {training_args.unfreeze_last_llm_layer}"
        )

    # [Explicit] Call to `freeze_backbones` here for clarity =>> will log exactly what is/is not frozen
    overwatch.info(f"Stage Info: ")
    model.freeze_backbones(stage)

    # Print number of total/trainable model parameters
    num_params = sum(p.numel() for p in model.parameters())
    num_trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    overwatch.info(
        f"# Parameters (in millions): {num_params / 10**6:.3f} Total, {num_trainable_params / 10**6:.3f} Trainable"
    )

    vla_dataset, action_tokenizer, collator = get_vla_dataset_and_collator(
        data_args.data_root_dir,
        data_args.data_mix,
        image_transform=model.vision_backbone.get_image_transform(),
        tokenizer=model.llm_backbone.get_tokenizer(),
        default_image_resolution=model.vision_backbone.default_image_resolution,