Function train_model

train_extrapolation.py:124–284

(model, dataloader, test_dataloader, args, test_type, start_epoch=0, start_time=time.time(),
                hop_count=None)

def train_model(model, dataloader, test_dataloader, args, test_type, start_epoch=0, start_time=time.time(),
                hop_count=None):
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay)
    total_training_steps = len(dataloader) * args.num_epochs
    scheduler = get_linear_schedule_with_warmup(
        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=total_training_steps
    )

    criterion = torch.nn.CrossEntropyLoss(label_smoothing=0.0)

    use_bf16 = args.precision == "bf16"
    scaler = torch.cuda.amp.GradScaler(enabled=not use_bf16)
    autocast_dtype = torch.bfloat16 if use_bf16 else torch.float16

    model.train()
    os.makedirs(args.checkpoint_dir, exist_ok=True)

    current_epoch = start_epoch
    while True:
        total_loss = 0.0
        progress_bar = tqdm(dataloader, desc=f"Epoch {current_epoch + 1}", unit="batch")

        for input_ids, target_tokens, attention_mask, input_lengths in progress_bar:

            if args.recurrence_type == 'dynamic':
                dynamic_rec = np.random.poisson(args.dynamic_mean)
                dynamic_rec = max(args.dynamic_min, min(args.dynamic_max, dynamic_rec))
                model.num_iterations = dynamic_rec
            else:
                model.num_iterations = args.recurrence

            input_ids = input_ids.to(args.device)
            target_tokens = target_tokens.to(args.device)
            attention_mask = attention_mask.to(args.device)

            optimizer.zero_grad()

            with torch.cuda.amp.autocast(dtype=autocast_dtype, enabled=True):
                outputs = model(input_ids=input_ids, attention_mask=attention_mask)
                if args.pred_pos == "last_token":
                    logits = outputs.logits[:, -1, :]
                else:
                    idx = (input_lengths - 1).view(-1, 1, 1).expand(-1, 1, outputs.logits.size(-1)).to(args.device)
                    logits = outputs.logits.gather(1, idx).squeeze(1)
                loss = criterion(logits, target_tokens)

            if use_bf16:
                loss.backward()
                if args.max_grad_norm > 0.0:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                optimizer.step()
            else:
                scaler.scale(loss).backward()
                if args.max_grad_norm > 0.0:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                scaler.step(optimizer)
                scaler.update()