Can't improve the scaler of batch size with ZeRO technique

[Lyn-Lucy]

Here is my code of training refer to the example of zero in [ColossalAI-Examples/train_v2.py at main · hpcaitech/ColossalAI-Examples (github.com)](https://github.com/hpcaitech/ColossalAI-Examples/blob/main/features/zero/train_v2.py)

def train(args, model1, kFold_num, kFold_idx):

    PLACEMENT_POLICY = 'auto'
    colossalai.launch_from_torch(config="config_colossal.py")
    pg = ProcessGroup()

    with ColoInitContext(device=get_current_device()):
        if "albert" in args.model_name_or_path:
            model = AlbertForClozeTest(args.model_name_or_path)
        elif "roberta" in args.model_name_or_path:
            model = RobertaForClozeTest(args.model_name_or_path)
        elif "bert" in args.model_name_or_path:
            model = BertForClozeTest(args.model_name_or_path)
        else:
            raise ValueError("model not supported")
    # print(list(model.parameters()))

    from colossalai.gemini import GeminiManager
    from colossalai.gemini.chunk import init_chunk_manager
    chunk_manager = init_chunk_manager(
        model=model,
        init_device=get_current_device(),
        search_range_mb=32
    )
    gemini_manager = GeminiManager(PLACEMENT_POLICY, chunk_manager)
    model = ZeroDDP(model, gemini_manager, pin_memory=True)
    model.to(args.device)
    
    
    


    """ Train the model """
    args.train_batch_size = args.per_gpu_train_batch_size * max(1, args.n_gpu)
    train_dataset = load_and_cache_examples(args, 'train', kFold_num=kFold_num, kFold_idx=kFold_idx)
    train_sampler = RandomSampler(train_dataset) if args.local_rank == -1 else DistributedSampler(train_dataset)
    train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=args.train_batch_size)


    test_dataloader  = train_dataloader

    # shard_strategy = TensorShardStrategy()
    # with ColoInitContext(device=torch.cuda.current_device()):
    #     model = AlbertForClozeTest(args.model_name_or_path).to(args.device)
    #     print(list(model.parameters()))
    #     print("==============================================")


    if args.max_steps > 0:
        t_total = args.max_steps
        args.num_train_epochs = args.max_steps // (len(train_dataloader) // args.gradient_accumulation_steps) + 1
    else:
        t_total = len(train_dataloader) //args.gradient_accumulation_steps * args.num_train_epochs

    eval_dataset = None
    if args.do_eval_during_train:
        eval_dataset = load_and_cache_examples(args, 'dev')

    # Prepare optimizer and schedule (linear warmup and decay)
    no_decay = ['bias', 'LayerNorm.weight']
    no_decay_params_id = [id(p) for n, p in model.named_parameters() if any(nd in n for nd in no_decay)]
    decay_params_id = [id(p) for n, p in model.named_parameters() if not any(nd in n for nd in no_decay)]
    optimizer_grouped_parameters = [
        {'params': [p for p in model.parameters() if id(p) in decay_params_id], 'weight_decay': args.weight_decay},
        {'params': [p for p in model.parameters() if id(p) in no_decay_params_id], 'weight_decay': 0.0}
    ]

    args.warmup_steps = int(t_total * args.warmup_rate)

    # optimizer = AdamW(optimizer_grouped_parameters, lr=args.learning_rate, eps=args.adam_epsilon)
    from colossalai.nn.optimizer import HybridAdam
    optimizer = HybridAdam(model.parameters(), lr=1e-3)
    optimizer = ZeroOptimizer(optimizer, model, initial_scale=2**5)

    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=args.warmup_steps,
                                                num_training_steps=t_total)
    scaler = torch.cuda.amp.GradScaler(enabled=bool(args.fp16))
    # scaler = torch.cuda.amp.GradScaler()

    loss_func = torch.nn.CrossEntropyLoss(reduction='none')
    


    # Train!
    logger.info("***** 训练阶段 *****")
    logger.info("  训练样本数 = %d", len(train_dataset))
    logger.info("  训练轮次 = %d", args.num_train_epochs)
    logger.info("  每GPU训练批次大小 = %d", args.per_gpu_train_batch_size)
    logger.info("  总训练批次大小 (w. parallel, distributed & accumulation) = %d",
                args.train_batch_size * args.gradient_accumulation_steps * (
                    dist.get_world_size() if args.local_rank != -1 else 1))
    logger.info("  梯度累积步数 = %d", args.gradient_accumulation_steps)
    logger.info("  总训练步数 = %d", t_total)

    fgm = None
    if args.adv_type == 'fgm':
        fgm = FGM(model)

    current_acc = 0.01
    global_step = 1
    tr_loss, logging_loss = 0.0, 0.0
    best_acc = 0

    optimizer.zero_grad()
    train_iterator = trange(int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0])
    set_seed(args)  # Added here for reproductibility (even between python 2 and 3)
    skip_scheduler = False
    for _ in train_iterator: # 遍历每个训练轮次
        epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])
        for step, batch in enumerate(epoch_iterator): # 遍历训练集
            model.train()
            batch = tuple(t.to(args.device) for t in batch)    
            with autocast(enabled=bool(args.fp16)):
            # if 1:
                article, option, answer, article_mask, option_mask, mask, blank_pos, sample_name = batch
                batch_size, option_num, out = model(batch)
                print("batch_size:",batch_size)
                target = answer.view(-1, )
                # calculate loss
                loss = loss_func(out, target)
                loss = loss.view(batch_size, option_num) * mask
                # replace nan to 0
                loss = torch.where(torch.isnan(loss), torch.full_like(loss, 0), loss)
                loss = loss.sum() / (mask.sum() if not mask.sum() == 0 else 1)

                if args.n_gpu > 1:
                   loss = loss.mean()  # mean() to average on multi-gpu parallel training
                if args.gradient_accumulation_steps > 1:
                   loss = loss / args.gradient_accumulation_steps

                scale_before_step = scaler.get_scale()
                # scaler.scale(loss).backward()
                optimizer.backward(loss)

                if args.adv_type == 'fgm':
                    fgm.attack()  # 对抗训练
                    out = model(batch)[2]
                    adv_loss = loss_func(out, target)
                    adv_loss = adv_loss.view(batch_size, option_num) * mask
                    adv_loss = torch.where(torch.isnan(adv_loss), torch.full_like(adv_loss, 0), adv_loss)
                    adv_loss = adv_loss.sum() / (mask.sum() if not mask.sum() == 0 else 1)
                    adv_loss.backward()
                    fgm.restore()

            tr_loss += loss.item()
            epoch_iterator.set_description("loss {}".format(round(loss.item()*args.gradient_accumulation_steps, 4)))
            if (step + 1) % args.gradient_accumulation_steps == 0: # 每满梯度累积步数则进行一次更新
                # scaler.unscale_(optimizer)
                # torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
                # scaler.step(optimizer)
                # scaler.update()
                if bool(args.fp16):
                    skip_scheduler = scaler.get_scale() != scale_before_step
                # if not skip_scheduler:
                #     scheduler.step()  # Update learning rate schedule
                before = list(model.parameters())
                # print(before)
                optimizer.step()
                after = list(model.parameters())
                # print(after)
                
                print("whether before parameters is same to the after :",before==after)
                optimizer.zero_grad()
                global_step += 1
                # 训练中验证
                if args.local_rank in [-1, 0] and args.logging_steps > 0 and args.do_eval_during_train and (
                        global_step % args.logging_steps == 0 or (global_step + 1) == t_total) and eval_dataset:
                    eval_result = evaluate(args, model,eval_dataset)
                    output_dir = args.output_dir
                    if not os.path.exists(output_dir):
                        os.makedirs(output_dir)
                    current_acc = eval_result['eval_accuracy']

                    logger.info("  最佳准确率 : {}".format(best_acc))
                    logger.info("  当前准确率 : {}".format(current_acc))
                    logger.info("  当前步数 : {}".format(global_step))
                    logger.info("  ")
                    for k in eval_result.keys():
                        logger.info("  {} : {}".format(k, eval_result[k]))
                    if current_acc > best_acc:
                        best_acc = current_acc
                        if not os.path.exists(args.output_dir) and args.local_rank in [-1, 0]:
                            os.makedirs(args.output_dir)
                        logger.info("保存模型到： %s", args.output_dir)
                        # model.save_pretrained(args.output_dir)
                        torch.save(model.state_dict(), f=args.output_dir)
                        output_eval_file = os.path.join(args.output_dir, "eval_results_during_train.txt")
                        with open(output_eval_file, "w") as writer:
                            for k, v in eval_result.items():
                                writer.write("{} : {}\n".format(k, v))

            if args.max_steps > 0 and global_step > args.max_steps:
                epoch_iterator.close()
                break
        if args.max_steps > 0 and global_step > args.max_steps:
            train_iterator.close()
            break
    return current_acc

However,I can only set one batch to per GPU. I can set 4 batch to per GPU without ZeRO.

Here is the result:

True 0.3443:  72%|███████▏  | 1330/1860 [17:53<07:34,  1.17it/s]
batch_size:1
whether before parameters is same to the after : True 

And the same time ,I also want to ask why that the parameters don't change after step.

[ver217]

Hi, ZeRO has own AMP. DO NOT use autocast and grad scaler.