Function train

ochat/training_deepspeed/train.py:203–308

()

def train():
    deepspeed.init_distributed(dist_backend="nccl")
    RANK = dist.get_rank()
    WORLD_SIZE = dist.get_world_size()

    # Args
    args       = parse_args()

    hub_upload_check(args.push_to_hub)

    # Dataset
    train_dataset, train_loader = create_dataset_and_dataloader(args, 0)

    if train_dataset is None:
        raise RuntimeError("Training data not found.")

    # Load model type
    args.model_type = train_dataset.metadata["model_type"]

    train_total_steps = args.epochs * train_dataset.estimate_num_batches()

    # Hyperparams
    args.lr = calculate_auto_lr(args.lr, args.batch_max_len, args.model_type, train_dataset)

    # Model
    model_engine, optimizer = create_model(args)

    # LR Scheduler
    lr_scheduler = create_lr_scheduler(args, train_total_steps)

    # Progress bar and logger
    progress_bar = None
    if RANK == 0:
        progress_bar = tqdm.tqdm(total=train_total_steps)

        wandb.init(project=args.wandb_project or os.path.basename(args.model_path), entity=args.wandb_entity, config=args)

    # Training Loop
    step = 0
    lr_this_step = None
    for epoch in range(args.epochs):
        print (f"[rank {RANK} of {WORLD_SIZE}]: Epoch {epoch}")

        ############ Load Dataset
        if epoch != 0:
            del train_dataset, train_loader

            train_dataset, train_loader = create_dataset_and_dataloader(args, epoch)

        ############ Train Epoch
        model_engine.train()
        for (batch_tensor, batch_info), all_numseq, cur_numseq in train_loader:
            step += 1
            if step > train_total_steps:  # At most train_total_steps
                break

            # To device
            batch_tensor = {k: (v.to(args.device) if v is not None else None) for k, v in batch_tensor.items()}