train.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#
import logging
import argparse
import getpass
import os
from time import time
import torch
from torch.utils.data import Subset
from src.metrics import (
    compute_accuracy,
    compute_accuracy_with_mask,
    compute_classification_metrics_per_cat,
)
from src.data import (
    parse_codecomplete_args,
    parse_goppa_args,
    parse_qc_args,
    get_datasets,
)
from src.model import get_model
from src.trainer import TrainingArguments, Trainer
from src.utils import end_wandb, init_wandb, initialize_exp, try_load_params, bool_flag


log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
logging.basicConfig(level=logging.INFO, format=log_format)


def get_params():
    parser = argparse.ArgumentParser(allow_abbrev=False)

    parser.add_argument("--seed", type=int, default=-1, help="-1 uses time() as seed")
    parser.add_argument("--resume", default="", help="Path to checkpoint .pt file")

    # Logging
    parser.add_argument("--log_every", type=int, default=100)
    parser.add_argument("--val_every", type=int, default=1000)
    parser.add_argument("--save_every", type=int, default=10_000)

    parser.add_argument("--data_path", type=str, required=False, default=None)
    parser.add_argument("--random_data_path", type=str, required=False)

    user = getpass.getuser()
    parser.add_argument("--dump_path", default=f"/checkpoint/{user}/dumped")
    parser.add_argument("--exp_name", default="debug_pretrain")
    parser.add_argument("--resume_from_checkpoint", default=None, type=str)

    # Model args
    parser.add_argument("--enc_emb_dim", type=int, default=512)
    parser.add_argument("--n_enc_layers", type=int, default=4)
    parser.add_argument("--n_enc_heads", type=int, default=8)
    parser.add_argument("--dropout", type=float, default=0)
    parser.add_argument("--attention_dropout", type=float, default=0)
    parser.add_argument(
        "--angular_emb",
        type=bool_flag,
        default=False,
        help="Whether to use xy coordinate embeddings",
    )

    parser.add_argument(
        "--compile", type=bool_flag, default=True, help="Use torch.compile?"
    )

    # Optimizer args
    parser.add_argument(
        "--optimizer",
        type=str,
        default="adam_warmup,lr=0.00001,warmup_updates=8000,weight_decay=0.001",
        help="Optimizer (SGD / RMSprop / Adam, etc.)",
    )
    parser.add_argument(
        "--timescale", type=int, default=40, help="How fast to decay the inv sqrt lr."
    )
    parser.add_argument(
        "--dtype", default="float32", choices=["float32", "float16", "bfloat16"]
    )

    # Training args
    parser.add_argument("--clip_grad_norm", type=float, default=5.0)
    parser.add_argument("--train_batch_size", type=int, default=256)
    parser.add_argument("--val_batch_size", type=int, default=512)
    parser.add_argument(
        "--eval_samples", type=int, default=1000, help="Number of evaluation samples"
    )
    parser.add_argument("--train_samples", type=int, default=2000000)
    parser.add_argument("--num_train_epochs", type=int, default=3)

    parser.add_argument("--shuffle", type=bool_flag, default=True)
    parser.add_argument("--workers", type=int, default=8, help="CPU workers for data")

    parser.add_argument(
        "--master_port", type=int, default=int(os.environ.get("MASTER_PORT", 10035))
    )
    parser.add_argument("--local_rank", type=int, default=-1)
    parser.add_argument("--device", type=str, default="cuda", help="Device to use")
    parser.add_argument("--is_master", type=bool_flag, default=True)
    parser.add_argument(
        "--multi_gpu", type=bool_flag, default=False, help="Run on multiple GPUs"
    )

    parser.add_argument("--task", type=str, default="lattice")
    parser.add_argument("--model", type=str, default="encoder")
    parser.add_argument("--Q", type=int, default=-1)

    parser.add_argument("--B", help="Angular Embedding Scale", type=int, default=1)
    parser.add_argument(
        "--K",
        help="Number of precision dimension in Angular embedding",
        type=int,
        default=1,
    )

    parser.add_argument("--max_hours", type=float, default=70, help="Max time allowed")
    parser.add_argument("--exp_id", type=str, default="", help="Experiment ID")
    parser.add_argument("--checkpoint_model", type=bool_flag, default=False)
    parser.add_argument("--wandb", type=bool_flag, default=False)
    parser.add_argument("--wandb_primary_key", type=str, default="exp_id")
    parser.add_argument("--tqdm", type=bool_flag, default=True)
    parser.add_argument("--copy_data", type=bool_flag, default=False)
    parser.add_argument("--tag", type=str, default="")
    parser.add_argument("--code_len", type=int, help="code length")
    parser.add_argument("--standard_only", type=bool_flag, default=True)
    parser.add_argument("--col_periods", type=str, default="")
    parser.add_argument("--row_periods", type=str, default="")

    params, unknown = parser.parse_known_args()

    if "qc" in params.task.split("-") or "mdpc" in params.task.split("-"):
        params = parse_qc_args(unknown, params)
    elif params.task.startswith("code-dist"):
        params = parse_goppa_args(unknown, params)
    elif params.task.startswith("code-complete"):
        params = parse_codecomplete_args(unknown, params)

    return params


def get_compute_metrics(params):
    if params.task.startswith("code-dist"):
        return compute_classification_metrics_per_cat
    elif params.task.startswith("code-ident"):
        return compute_accuracy
    else:
        return compute_accuracy_with_mask


if __name__ == "__main__":
    os.environ["CUDA_LAUNCH_BLOCKING"] = str(1)

    params = get_params()

    if params.seed < 0:
        params.seed = int(time()) % 1000000

    try_load_params(params)

    logger = initialize_exp(params)

    train_dataset, test_dataset = get_datasets(params)

    model = get_model(params)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"Total parameters: {total_params}")

    report_to = init_wandb(params)

    local = params.local_rank == -1
    logger.info(f"local rank is {params.local_rank} out of {torch.cuda.device_count()}")

    training_args = TrainingArguments(
        dump_path=params.dump_path,
        evaluation_strategy="steps",
        num_train_epochs=params.num_train_epochs,
        eval_steps=params.val_every,
        logging_steps=params.log_every,
        save_steps=params.save_every,
        per_device_train_batch_size=params.train_batch_size,
        per_device_eval_batch_size=params.val_batch_size,
        report_to=report_to,
        local_rank=params.local_rank,
        dataloader_num_workers=params.workers,
        device=params.device,
        multi_gpu=params.multi_gpu,
        dtype=params.dtype,
        max_grad_norm=params.clip_grad_norm,
        compile=params.compile,
        optimizer=params.optimizer,
        resume_from_checkpoint=params.resume_from_checkpoint,
    )

    callbacks = []

    trainer = Trainer(
        model=model,
        training_args=training_args,
        args=params,
        train_dataset=train_dataset,
        eval_dataset=test_dataset,
        data_collator=(
            train_dataset.dataset.collate_fn
            if isinstance(train_dataset, Subset)
            else train_dataset.collate_fn
        ),
        compute_metrics=get_compute_metrics(params),
        callbacks=callbacks,
    )

    logger.info(f"gpu count {torch.cuda.device_count()}")

    if params.num_train_epochs > 0:
        trainer.train()
    else:
        trainer.evaluate()

    end_wandb(params)