Repository: hpcaitech/ColossalAI
File: examples/language/bert/finetune.py (Python)

Code before:
                                      precision='fp16',
                                      initial_scale=1)

    booster = Booster(plugin=plugin, **booster_kwargs)

    # ==============================
    # Prepare Dataloader
    # ==============================
    data_builder = GLUEDataBuilder(model_name,
                                   plugin,
                                   args.task,
                                   train_batch_size=BATCH_SIZE,
                                   eval_batch_size=BATCH_SIZE)
    train_dataloader = data_builder.train_dataloader()
    test_dataloader = data_builder.test_dataloader()

    # ====================================
    # Prepare model, optimizer
    # ====================================
    # bert pretrained model

    cfg = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels)

    # lazy_init
    use_lazy_init = args.use_lazy_init
    ctx = LazyInitContext() if use_lazy_init else nullcontext()
    with ctx:
        if model_name == "bert-base-uncased":
            model = BertForSequenceClassification.from_pretrained(model_name, config=cfg).cuda()
        elif model_name == "albert-xxlarge-v2":
            model = AlbertForSequenceClassification.from_pretrained(model_name, config=cfg)
        else:
            raise RuntimeError

    # optimizer
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],
            "weight_decay": WEIGHT_DECAY,
        },
        {
            "params": [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay)],
            "weight_decay": 0.0,
        },
    ]

    #TODO something wrong with HybridAdam when using pp
    # optimizer = HybridAdam(optimizer_grouped_parameters, lr=lr, eps=1e-8)
    optimizer = Adam(model.parameters(), lr=1e-3)


Review comment (suggestion):
Lazy init context does not support `from_pretrained()` now. It's not essential in this example. You can remove it.

Code after:
        booster_kwargs['mixed_precision'] = 'fp16'
    if args.plugin.startswith('torch_ddp'):
        plugin = TorchDDPPlugin()
    elif args.plugin == 'gemini':
        plugin = GeminiPlugin(placement_policy='cuda', strict_ddp_mode=True, initial_scale=2**5)
    elif args.plugin == 'low_level_zero':
        plugin = LowLevelZeroPlugin(initial_scale=2**5)
    elif args.plugin == 'hybrid_parallel':

        # modify the param accordingly for finetuning test cases
        plugin = HybridParallelPlugin(tp_size=1,
                                      pp_size=2,
                                      num_microbatches=None,
                                      microbatch_size=1,
                                      enable_all_optimization=True,
                                      zero_stage=1,
                                      precision='fp16',
                                      initial_scale=1)

    booster = Booster(plugin=plugin, **booster_kwargs)

    # ==============================
    # Prepare Dataloader
    # ==============================
    data_builder = GLUEDataBuilder(model_name,
                                   plugin,
                                   args.task,
                                   train_batch_size=BATCH_SIZE,
                                   eval_batch_size=BATCH_SIZE)
    train_dataloader = data_builder.train_dataloader()
    test_dataloader = data_builder.test_dataloader()

    # ====================================
    # Prepare model, optimizer
    # ====================================
    # bert pretrained model

    cfg = AutoConfig.from_pretrained(model_name, num_labels=data_builder.num_labels)

    if model_name == "bert-base-uncased":
        model = BertForSequenceClassification.from_pretrained(model_name, config=cfg).cuda()
    elif model_name == "albert-xxlarge-v2":
        model = AlbertForSequenceClassification.from_pretrained(model_name, config=cfg)
    else:
        raise RuntimeError

    # optimizer
    no_decay = ["bias", "LayerNorm.weight"]
    optimizer_grouped_parameters = [
        {
            "params": [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)],