Save/load weight when using ZeRO

[Zjh-819]

Hello, I don't know how to save model's params properly when using ZeRO. I used the save_checkpoint API, but when I check the .pt file manually, I found all params are of value '[]' except some biases. It seems that another user has encounetered the same problem, too(#1082 (comment)). But I didn't find any solution now. How can I fix this?
environment: python 3.8.13 + pytorch 1.10.1 + cu113 + ColossalAI 0.18 + huggingface transformer 4.20.1
Here are my code and config.

import torch as th
from dataset import get_examples, GSMDataset
from transformers import GPT2Tokenizer, GPT2LMHeadModel
from transformers import GPT2Config
from transformers import get_scheduler
from tqdm import tqdm
from torch.utils.data import DataLoader

import colossalai
from colossalai.core import global_context as gpc
from colossalai.logging import get_dist_logger
from colossalai.zero.init_ctx import ZeroInitContext
from colossalai.nn.optimizer import HybridAdam
from colossalai.utils import save_checkpoint, load_checkpoint


class CWrapper:
    def __init__(self):
        super().__init__()

    def forward(self, logits, labels):
        return logits


def main():
    # parse arguments
    args = colossalai.get_default_parser().parse_args()

    # launch distributed environment
    colossalai.launch(config='./colossal_config.py',
                      rank=args.rank,
                      world_size=args.world_size,
                      host=args.host,
                      port=args.port,
                      backend=args.backend
                      )

    # torch launcher
    # colossalai.launch_from_torch(config='./colossal_config.py')

    logger = get_dist_logger()
    model_name = 'gpt2-xl'
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    train_examples = get_examples("train_no_annotation")[:10]
    train_dset = GSMDataset(tokenizer, train_examples, loss_on_prefix=True)

    device = th.device("cuda")
    config = GPT2Config.from_pretrained(model_name)

    # ZeRO init
    with ZeroInitContext(target_device=th.device('cuda', 0),
                         shard_strategy=gpc.config.zero.model_config.shard_strategy,
                         shard_param=False):
        model = GPT2LMHeadModel.from_pretrained(model_name, config=config)
        # load_checkpoint(file="ckpts/model_ckpt_gpt2-xl.pt", model=model)  # 必须放在同一个context之下，否则报错

    optim = HybridAdam(model.parameters(), lr=1e-5)
    train_loader = DataLoader(train_dset, batch_size=2, shuffle=True, pin_memory=True)

    num_epochs = 10
    num_training_steps = num_epochs * len(train_loader)
    lr_scheduler = get_scheduler(
        "linear",
        optimizer=optim,
        num_warmup_steps=0,
        num_training_steps=num_training_steps,
    )

    engine, train_loader, _, lr_scheduler = colossalai.initialize(model=model,
                                                                  optimizer=optim,
                                                                  train_dataloader=train_loader,
                                                                  lr_scheduler=lr_scheduler
                                                                  )

    # logger.info('load success!')
    engine.train()
    for epoch in range(10):
        for batch in tqdm(train_loader):
            engine.zero_grad()
            batch = {k: v.to(device) for k, v in batch.items()}
            outputs = engine(**batch, labels=batch["input_ids"])
            loss = outputs.loss
            engine.backward(loss)
            engine.step()
            lr_scheduler.step()

        save_checkpoint(file="ckpts/model_ckpt_" + model_name + ".pt", epoch=epoch, model=model)
        logger.info(f'epoch: {epoch}')
    save_checkpoint(file="ckpts/model_ckpt_" + model_name + ".pt", epoch=epoch, model=model)

    # model.save_pretrained("model_ckpts_" + model_name + "/")
    # th.save(model.state_dict(), "ckpts/model_ckpt_" + model_name + str(epoch) + ".pt")


if __name__ == "__main__":
    main()

from colossalai.zero.shard_utils import TensorShardStrategy

BATCH_SIZE = 2
NUM_EPOCHS = 60
SEQ_LEN = 1024

zero = dict(
    model_config=dict(
        shard_strategy=TensorShardStrategy(),
        reduce_scatter_bucket_size_mb=25,
        tensor_placement_policy="auto",
        reuse_fp16_shard=True
    ),
    optimizer_config=dict(
        gpu_margin_mem_ratio=0.72,
        # initial_scale=2 ** 5,
        # min_scale=1,
        # growth_factor=2,
        # backoff_factor=0.5,
        # growth_interval=1000,
        # hysteresis=2,
        # max_scale=2 ** 32
    )
)

[ver217]

Hi, you should save engine.model

[Zjh-819]

Thanks, it works!