Function setup_model_and_optimizer

train_utils.py:243–267

Setup model and optimizer.

(args, model_type=None, multi_token=True, num_labels=None, spell_length=None)

def setup_model_and_optimizer(args, model_type=None, multi_token=True, num_labels=None, spell_length=None):
    """Setup model and optimizer."""

    model = get_model(args, model_type=model_type, multi_token=multi_token, num_labels=num_labels,
                      spell_length=spell_length)
    param_groups = get_optimizer_param_groups(model)

    if args.train_data is not None or args.data_dir is not None and (args.epochs > 0 or args.train_iters > 0):
        if args.deepspeed:
            print_rank_0("DeepSpeed is enabled.")

            model, optimizer, _, _ = deepspeed.initialize(
                model=model,
                model_parameters=param_groups,
                args=args,
                mpu=mpu,
                dist_init_required=False
            )
        else:
            optimizer = get_optimizer(param_groups, args)
        lr_scheduler = get_learning_rate_scheduler(optimizer, args)
    else:
        optimizer, lr_scheduler = None, None

    return model, optimizer, lr_scheduler


def backward_step(optimizer, model, lm_loss, args, timers):