halo，I continue pretrain llama2-13B model ，but save state_dict is about 50GB file
Labels: oncall: distributed, triaged

### 🐛 Describe the bug

@record
def training_function(args):
    # get some base rank info
    # metric = evaluate.load("glue", "mrpc")
    world_size = os.getenv("WORLD_SIZE")
    rank = os.getenv("RANK")
    local_rank = os.getenv("LOCAL_RANK")
    assert world_size is not None, f"WORLD_SIZE is needed {world_size}"
    assert rank is not None, f"RANK is needed {rank}"
    assert local_rank is not None, f"RANK is needed {local_rank}"
    world_size = int(world_size)
    rank = int(rank)
    local_rank = int(local_rank)

    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
    # Load the pre-trained model and setup its configuration
    model = LlamaForCausalLM.from_pretrained(
            args.model_name_and_path,
            load_in_8bit=True if args.quantization else None,
            device_map="auto" if args.quantization else None,
            return_dict=True
        )
    model.train()
    # Initialize accelerator
    gradient_accumulation_steps = args.gradient_accumulation_steps
    if args.with_tracking:
        accelerator = Accelerator(
            gradient_accumulation_steps=gradient_accumulation_steps, mixed_precision=args.mixed_precision, log_with="tensorboard", project_dir=args.project_dir)
    else:
        accelerator = Accelerator(gradient_accumulation_steps=gradient_accumulation_steps, mixed_precision=args.mixed_precision)

        # We need to initialize the trackers we use, and also store our configuration
    if args.with_tracking:
        run = os.path.split(__file__)[-1].split(".")[0]
        tracker_cfg = {
            'lr': args.lr,
            'per_device_batch_size': args.batch_size,
            'seed': args.seed,
            'num_epoch': args.num_epochs
        }
        accelerator.init_trackers(run, tracker_cfg)
    accelerator.print(f"***** Total GPUS: {accelerator.num_processes} *****")
    # set seed
    set_seed(args.seed)

    
    # epoch or steps
    if hasattr(args.checkpointing_steps, "isdigit"):
        if args.checkpointing_steps == "epoch":
            checkpointing_steps = args.checkpointing_steps
        elif args.checkpointing_steps.isdigit():
            checkpointing_steps = int(args.checkpointing_steps)
        else:
            raise ValueError(
                f"Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed."
            )
    else:
        checkpointing_steps = None

    if hasattr(args.lora_save_steps, "isdigit"):
        if args.lora_save_steps == "epoch":
            lora_save_steps = args.lora_save_steps
        elif args.lora_save_steps.isdigit():
            lora_save_steps = int(args.lora_save_steps)
        else:
            raise ValueError(
                f"Argument `lora_save_steps` must be either a number or `epoch`. `{args.lora_save_steps}` passed."
            )
    else:
        lora_save_steps = None

    
    # Load the tokenizer and add special tokens
    if args.token_name:
        
        accelerator.print(f'Use tokenizer: [{args.token_name}]')
        tokenizer = LlamaTokenizer.from_pretrained(args.token_name, )
        model.config.vocab_size = tokenizer.vocab_size
        model.config.pad_token_id = tokenizer.pad_token_id
        model.config.bos_token_id = tokenizer.bos_token_id
        model.config.eos_token_id = tokenizer.eos_token_id
        model.resize_token_embeddings(tokenizer.vocab_size)
    else:
        tokenizer = LlamaTokenizer.from_pretrained(args.model_name_and_path)
        tokenizer.add_special_tokens(
                {
                
                    "pad_token": "<PAD>",
                }
            )
        
    # get processed dataset
    # vol = os.getenv('LPAI_INPUT_DATA_0')
    # dataset_all = load_from_disk(f'{vol}/NLP/bigDatasets/pretrain_D/tokened_grouped_eval')
    dataset_all = load_from_disk(args.dataset_path)
    dataset_all = dataset_all.train_test_split(test_size=0.01)
    dataset_train = dataset_all['train']
    dataset_val = dataset_all['test']
    train_sampler = None
    val_sampler = None
    # print(f'dist.get_rank() ->{dist.get_rank()}, dist.get_world_size() -> {dist.get_world_size()}')
    
    def collate_fn(examples):
        # On TPU it's best to pad everything to the same length or training will be very slow.
        max_length = 128 if accelerator.distributed_type == DistributedType.TPU else None
        # When using mixed precision we want round multiples of 8/16
        if accelerator.mixed_precision == "fp8":
            pad_to_multiple_of = 16
        elif accelerator.mixed_precision != "no":
            pad_to_multiple_of = 8
        else:
            pad_to_multiple_of = None
        
        return tokenizer.pad(
            examples,
            padding=True,
            max_length=max_length,
            pad_to_multiple_of=pad_to_multiple_of,
            return_tensors="pt",
        )
    
    if args.enable_fsdp:
        train_sampler = DistributedSampler(
            dataset_train,
            rank=dist.get_rank(),
            # rank=accelerator.device,
            seed=args.seed,
            num_replicas=dist.get_world_size(),
            shuffle=True,
        )
        if args.run_validation:
            val_sampler = DistributedSampler(
                dataset_val,
                rank=dist.get_rank(),
                # rank=accelerator.device,
                num_replicas=dist.get_world_size(),
            )
    
    train_dataloader = torch.utils.data.DataLoader(
        dataset_train,
        batch_size=args.batch_size,
        num_workers=args.num_workers_dataloader,
        pin_memory=True,
        sampler=train_sampler if train_sampler else None,
        drop_last=True,
        # collate_fn=default_data_collator,
        collate_fn=collate_fn,
    )
    
    if args.run_validation:
        eval_dataloader = torch.utils.data.DataLoader(
            dataset_val,
            batch_size=args.val_batch_size,
            num_workers=args.num_workers_dataloader,
            pin_memory=True,
            sampler=val_sampler if val_sampler else None,
            drop_last=True,
            # collate_fn=default_data_collator,
            collate_fn=collate_fn,
        )    
    
    if args.use_peft:
        model.to(torch.bfloat16)
        if args.resume_from_lora:
            assert args.resume_from_checkpoint is None, 'lora adapter or checkpoint just need use any one'
            if args.resume_from_lora is not None or args.resume_from_lora != "":
                accelerator.print(f"Resumeing from lora_adapter: {args.resume_from_lora}")
                model=PeftModel.from_pretrained(model, args.resume_from_lora, is_trainable=True)
                model.print_trainable_parameters()
            else:
            # Get the most recent checkpoint
                dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
                dirs.sort(key=os.path.getctime)
                path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
                accelerator.print(f"Resumeing from lora_adapter: {path}")
                model=PeftModel.from_pretrained(model, path, is_trainable=True)
                model.print_trainable_parameters()
        else:
        # target_modules = args.lora_modules.split(',')
        # accelerator.print(target_modules)
            # model.to(torch.bfloat16)
            peft_config = LoraConfig(
                    r=args.lora_r,
                    lora_alpha=32,
                    # target_modules=target_modules,
                    target_modules=["q_proj", "v_proj", "o_proj", "k_proj", "gate_proj", "up_proj"],
                    bias="none",
                    task_type="CAUSAL_LM",
                    lora_dropout=0.05,
                    inference_mode=False
            )
            
            model = get_peft_model(model, peft_config)
            model.print_trainable_parameters()
    
    
    mixed_precision_policy, wrapping_policy = get_policies(fsdp_config, rank)
    my_auto_wrapping_policy = fsdp_auto_wrap_policy(model, LlamaDecoderLayer)
    bfSixteen = MixedPrecision(
            param_dtype=torch.bfloat16,
            # Gradient communication precision.
            reduce_dtype=torch.bfloat16,
            # Buffer precision.
            buffer_dtype=torch.bfloat16,
            cast_forward_inputs=True,
        )
    torch.cuda.set_device(local_rank)
    model = FSDP(
        model,
        auto_wrap_policy= my_auto_wrapping_policy if args.use_peft else wrapping_policy,
        # mixed_precision=mixed_precision_policy if not fsdp_config.pure_bf16 else None,
        mixed_precision=bfSixteen,
        sharding_strategy=fsdp_config.sharding_strategy,
        backward_prefetch=BackwardPrefetch.BACKWARD_PRE,
        device_id=torch.cuda.current_device(),
        # forward_prefetch=True,
        limit_all_gathers=True,
    )
    if fsdp_config.fsdp_activation_checkpointing:
        policies.apply_fsdp_checkpointing(model)
    # We need to keep track of how many total steps we have iterated over
    overall_step = 0
    # We also need to keep track of the stating epoch so files are named properly
    starting_epoch = 0
    
    optimizer = AdamW(params=model.parameters(), lr=args.lr)
    # Instantiate scheduler
    lr_scheduler = get_linear_schedule_with_warmup(
        optimizer=optimizer,
        num_warmup_steps=500,
        num_training_steps=(len(train_dataloader) * args.num_epochs) // gradient_accumulation_steps,
    )

    # Prepare everything
    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
    # prepare method.
    optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
        optimizer, train_dataloader, eval_dataloader, lr_scheduler
    )

    total_batch_size = args.batch_size * accelerator.num_processes * args.gradient_accumulation_steps
    num_training_steps=int((len(train_dataloader) * args.num_epochs) // gradient_accumulation_steps)
    accelerator.print("***** Running training *****")
    accelerator.print(f"  Num examples = {len(dataset_train)}")
    accelerator.print(f"  Num train_loader = {len(train_dataloader)}")
    accelerator.print(f"  Num eval_loader = {len(eval_dataloader)}")
    accelerator.print(f"  Num Epochs = {args.num_epochs}")
    accelerator.print(f"  Instantaneous batch size per device = {args.batch_size}")
    accelerator.print(f"  Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}")
    accelerator.print(f"  Gradient Accumulation steps = {args.gradient_accumulation_steps}")
    accelerator.print(f"  Total optimization steps = {num_training_steps}")
    progress_bar = tqdm(range(num_training_steps), disable=not accelerator.is_local_main_process)
    
    
    # Potentially load in the weights and states from a previous save
    if args.resume_from_lora:
        assert args.resume_from_checkpoint is None, 'lora adapter or checkpoint just need use any one'
        if args.resume_from_lora is not None or args.resume_from_lora != "":
            path = os.path.basename(args.resume_from_lora)
        else:
            # Get the most recent checkpoint
            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
            dirs.sort(key=os.path.getctime)
            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
        # Extract `epoch_{i}` or `step_{i}`
        training_difference = os.path.splitext(path)[0]
        if "epoch" in training_difference:
            starting_epoch = int(training_difference.replace("lora_epoch_", "")) + 1
            resume_step = None
            resume_step_r = starting_epoch * len(train_dataloader)
        else:
            resume_step_r = int(training_difference.replace("lora_step_", ""))
            starting_epoch = resume_step_r // len(train_dataloader)
            resume_step = resume_step_r - (starting_epoch * len(train_dataloader))
            accelerator.print(f'***** In use_peft resume_step {resume_step}')
    
    if args.resume_from_checkpoint:
        if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
            accelerator.print(f"***** Resumed from checkpoint: {args.resume_from_checkpoint} *****")
            accelerator.load_state(args.resume_from_checkpoint)
            path = os.path.basename(args.resume_from_checkpoint)
        else:
            # Get the most recent checkpoint
            dirs = [f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
            dirs.sort(key=os.path.getctime)
            path = dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
            accelerator.print(f"***** Resumed from checkpoint: {path} *****")
            accelerator.load_state(args.resume_from_checkpoint)
        # Extract `epoch_{i}` or `step_{i}`
        training_difference = os.path.splitext(path)[0]

        if "epoch" in training_difference:
            starting_epoch = int(training_difference.replace("epoch_", "")) + 1
            resume_step = None
            resume_step_r = starting_epoch * len(train_dataloader)
        else:
            resume_step_r = int(training_difference.replace("step_", ""))
            starting_epoch = resume_step_r // len(train_dataloader)
            resume_step = resume_step_r - (starting_epoch * len(train_dataloader))
            accelerator.print(f'***** In checkpoint resume_step {resume_step}')
            # overall_step = resume_step // args.gradient_accumulation_steps
    # update the progress_bar if load from checkpoint
    
    
        # Create a gradient scaler for fp16
    if train_config.use_fp16 and train_config.enable_fsdp:
        scaler = ShardedGradScaler()
    elif train_config.use_fp16 and not train_config.enable_fsdp:
        scaler = torch.cuda.amp.GradScaler() 
    if train_config.enable_fsdp:
        world_size = int(os.environ["WORLD_SIZE"]) 
    
    
    # Now we train the model
    for epoch in range(starting_epoch, args.num_epochs):
        model.train()
        train_sampler.set_epoch(epoch)
        accelerator.print(f'***** Ins loop epoch and distribute set epoch *****')
        if args.with_tracking:
            total_loss = 0
        if (args.resume_from_checkpoint or args.resume_from_lora) and epoch == starting_epoch and resume_step is not None:
            # We need to skip steps until we reach the resumed step
            accelerator.print(f'accelerator.skip_first_batches -> resume_step {resume_step}')
            # active_dataloader = accelerator.skip_first_batches(train_dataloader, resume_step)
            active_dataloader = skip_first_batches(train_dataloader, resume_step)
            overall_step += resume_step_r
            progress_bar.update(overall_step)
            accelerator.print(f'***** oversteps {overall_step} skip train loader *****')
        else:
            # After the first iteration though, we need to go back to the original dataloader
            active_dataloader = train_dataloader
        first_f = False
        for step, batch in enumerate(active_dataloader):
            # We could avoid this line since we set the accelerator with `device_placement=True`.
            # batch.to(accelerator.device)
            with accelerator.accumulate(model):
                if not first_f:
                    log_ids = batch['input_ids']
                    accelerator.print(tokenizer.batch_decode(log_ids))
                    first_f = True
                for key in batch.keys():
                    if args.enable_fsdp:
                        batch[key] = batch[key].to(accelerator.device)
                        # batch[key] = batch[key].to(local_rank)
                    else:
                        batch[key] = batch[key].to('cuda:0')   
                loss = model(**batch).loss
                # loss = loss / gradient_accumulation_steps
            # We keep track of the loss at each epoch
            if args.with_tracking:
                total_loss += loss.detach().float()
            accelerator.backward(loss)
            accelerator.log({"loss": loss.item()}, step=step)
            accelerator.print(f"step: {overall_step} loss: {loss.item()}")
            if accelerator.sync_gradients:
                accelerator.clip_grad_norm_(model.parameters(), 1.0)
            # if step % gradient_accumulation_steps == 0:
                optimizer.step()
                lr_scheduler.step()
                optimizer.zero_grad()
            if accelerator.sync_gradients:
                overall_step += 1
                progress_bar.update(1)

            if isinstance(checkpointing_steps, int):
                output_dir = f"step_{overall_step}"
                if overall_step % checkpointing_steps == 0:
                    if args.output_dir is not None:
                        output_dir = os.path.join(args.output_dir, output_dir)
                    accelerator.save_state(output_dir)
                    
                    # 0.0
                    # keep_fp32_wrapper
                    # create singleton saving policies to avoid making over and over
                    fullstate_save_policy = FullStateDictConfig(offload_to_cpu=True, rank0_only=True)
                    with FSDP.state_dict_type(
                            model, StateDictType.FULL_STATE_DICT, fullstate_save_policy
                        ):
                        accelerator.print(f'***** model.state_dict() {model.state_dict().keys()} *****')
                        cpu_state = model.state_dict()
                        # pwd_dir = Path('./').absolute() / args.output_dir
                        pwd_dir = os.path.join(args.output_dir, f'epoch_model.pth')
                        # torch.save(cpu_state, args.output_dir)
                        # accelerator.print(f'***** pwd dir: {pwd_dir} *****')
                        accelerator.save(cpu_state, pwd_dir)

### Versions

Collecting environment information...
PyTorch version: 2.0.1+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.26.0
Libc version: glibc-2.27

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-3.10.0-1160.42.2.el7.x86_64-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 11.7.99
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB
GPU 1: NVIDIA A100-SXM4-40GB
GPU 2: NVIDIA A100-SXM4-40GB
GPU 3: NVIDIA A100-SXM4-40GB
GPU 4: NVIDIA A100-SXM4-40GB
GPU 5: NVIDIA A100-SXM4-40GB
GPU 6: NVIDIA A100-SXM4-40GB
GPU 7: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.4
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.4
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
架构：           x86_64
CPU 运行模式：   32-bit, 64-bit
字节序：         Little Endian
CPU:             112
在线 CPU 列表：  0-111
每个核的线程数： 2
每个座的核数：   28
座：             2
NUMA 节点：      2
厂商 ID：        GenuineIntel
CPU 系列：       6
型号：           106
型号名称：       Intel(R) Xeon(R) Platinum 8350C CPU @ 2.60GHz
步进：           6
CPU MHz：        2593.904
BogoMIPS：       5187.80
超管理器厂商：   KVM
虚拟化类型：     完全
L1d 缓存：       48K
L1i 缓存：       32K
L2 缓存：        1280K
L3 缓存：        49152K
NUMA 节点0 CPU： 0-55
NUMA 节点1 CPU： 56-111
标记：           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl xtopology nonstop_tsc eagerfpu pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb ibrs_enhanced fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq spec_ctrl arch_capabilities

Versions of relevant libraries:
[pip3] flake8==3.8.4
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.21.5
[pip3] pytorch-lightning==1.6.5
[pip3] pytorch-metric-learning==1.5.2
[pip3] pytorch-ranger==0.1.1
[pip3] pytorch-wpe==0.0.1
[pip3] torch==2.0.1
[pip3] torch-audiomentations==0.11.0
[pip3] torch-complex==0.4.3
[pip3] torch-optimizer==0.1.0
[pip3] torch-pitch-shift==1.2.2
[pip3] torch-stoi==0.1.2
[pip3] torchaudio==2.0.2
[pip3] torchmetrics==0.7.2
[pip3] torchvision==0.15.2
[pip3] triton==2.0.0
[conda] numpy                     1.21.5                   pypi_0    pypi
[conda] pytorch-lightning         1.6.5                    pypi_0    pypi
[conda] pytorch-metric-learning   1.5.2                    pypi_0    pypi
[conda] pytorch-ranger            0.1.1                    pypi_0    pypi
[conda] pytorch-wpe               0.0.1                    pypi_0    pypi
[conda] torch                     2.0.1                    pypi_0    pypi
[conda] torch-audiomentations     0.11.0                   pypi_0    pypi
[conda] torch-complex             0.4.3                    pypi_0    pypi
[conda] torch-optimizer           0.1.0                    pypi_0    pypi
[conda] torch-pitch-shift         1.2.2                    pypi_0    pypi
[conda] torch-stoi                0.1.2                    pypi_0    pypi
[conda] torchaudio                2.0.2                    pypi_0    pypi
[conda] torchmetrics              0.7.2                    pypi_0    pypi
[conda] torchvision               0.15.2                   pypi_0    pypi
[conda] triton                    2.0.0                    pypi_0    pypi

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu