microsoft / DeepSpeedExamples

Example models using DeepSpeed
Apache License 2.0
5.98k stars 1.02k forks source link

Backward through the graph twice in training #777

Closed MIkumikumi0116 closed 11 months ago

MIkumikumi0116 commented 11 months ago

Describe the bug When I am training my model, everything goes well in the first batch, but an error occurs in the second batch: RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

To Reproduce I am fine-tuning llama2 using my own implementation of LoRA. Each lora linear in my llama2 has different LoRA rank. Each linear in my model is replaced to

class Adaptive_Linear(torch.nn.Module):
    def __init__(
        self,
        model_arg: dict, # batch size, max seqlen, vocab size, etc.
        name: str, # example: layer_list.0.attention.key
        lora_rank_dict: dict, # lora rank of each linear
        in_features: int,
        out_features: int 
    ):
        super().__init__()

        self.name      = name
        self.model_arg = model_arg

        self.linear    = torch.nn.Linear(in_features, out_features, bias = False)

        self.lora_rank = lora_rank_dict[self.name]
        if self.lora_rank > 0:
            self.lora_dropout  = torch.nn.Dropout(model_arg.lora_dropout)
            self.lora_a_linear = torch.nn.Linear(in_features, self.lora_rank,  bias = False)
            self.lora_b_linear = torch.nn.Linear(self.lora_rank, out_features, bias = False)

            self.lora_a_linear.weight.data = torch.rand_like(self.lora_a_linear.weight.data) - 1 * 2
            self.lora_b_linear.weight.data = torch.zeros_like(self.lora_b_linear.weight.data)

    def trainable_parameters(self):
        # only train lora matrix
        # the trainable parameters of my model are trainable_parameters
        # of all Adaptive_Linear connected together using itertools.chain.
        return iter((self.lora_a_linear.weight, self.lora_b_linear.weight))

    def forward(self, x):
        hidden = self.linear(x)

        if self.lora_rank > 0:
            input  = self.lora_dropout(x)
            lora_a = self.lora_a_linear(input)
            lora_b = self.lora_b_linear(lora_a)
            output = hidden + lora_b
        else:
            output = hidden

        return output

My dataset is

class Dataset(torch.utils.data.Dataset):
  def __init__(self, model_arg, dataset_path):
      self.model_arg    = model_arg

      self.dataset_path = dataset_path
      self.dataset      = torch.load(self.dataset_path) # List[Tuple] list of tuples of (input token list, output token list, label) token and label are int

      self.dataset      = [data_item for data_item in self.dataset if len(data_item[0]) >= MIN_SEQ_LEN and len(data_item[0]) <= MAX_SEQ_LEN]

  def __len__(self):
      return len(self.dataset)

  def __getitem__(self, index):
      input_token, output_token, label = self.dataset

      return input_token, output_token, label

Collate function to pad tokens of a batch to same length

collate_fn(batch):
    input_token_list, output_token_list, label_list = zip(*batch)

    padded_input_token  = torch.nn.utils.rnn.pad_sequence([seq for seq in input_token_list],  batch_first = True)
    padded_output_token = torch.nn.utils.rnn.pad_sequence([seq for seq in output_token_list], batch_first = True)

    return padded_input_token, padded_output_token, label_list

Initialize model and deepspeed

model = Llama_2_With_Adaptive_Lora(model_arg, lora_rank_dict)
model.load_state_dict(state_dict)
for name, parameter in model.named_parameters():
    if not ('lora' in name):
        parameter.requires_grad = False

engine, optimizer, train_dataloader, _= deepspeed.initialize(
    model            = model,
    model_parameters = model.trainable_parameters(),
    training_data    = train_dataset,
    args             = args, # local_rank only
    collate_fn       = collate_fn,
    config           = deepspeed_config
)

Deepspeed config

{
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 8,
    "gradient_clipping": 1.0,
    "steps_per_print": 10,
    "wall_clock_breakdown": false,
    "memory_breakdown": false,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [
                0.9,
                0.95
            ],
            "eps": 1e-8,
            "weight_decay": 1e-4
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 2e-5,
            "warmup_num_steps": 1000,
            "total_num_steps": 10000
        }
    },
    "bf16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },

        "allgather_partitions": true,
        "allgather_bucket_size": 1e7,
        "reduce_scatter": true,
        "reduce_bucket_size": 1e7,

        "overlap_comm": true,
        "contiguous_gradients": true,

        "sub_group_size": 1e9,
        "stage3_prefetch_bucket_size": 1e7,
        "stage3_param_persistence_threshold": 1e5,
        "stage3_max_live_parameters": 1e8,
        "stage3_max_reuse_distance": 1e8,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "activation_checkpointing": {
        "partition_activations": true,
        "number_checkpoints": 100,

        "cpu_checkpointing": true,
        "contiguous_memory_optimization": true,
        "synchronize_checkpoint_boundary": false,
        "profile": true
    },
    "data_types": {
        "grad_accum_dtype":"bf16"
    }
}

Main code of training

engine.train()
for input_token, target_token, _ in train_dataloader:
    input_token  = input_token. to(engine.device)
    target_token = target_token.to(engine.device)

    pred_logits  = engine(input_token)

    target_pos   = (target_token != PAD_ID) # only calculate loss on non-pad token
    pred_logits  = pred_logits[ target_pos].view(-1, model_arg.vocab_size)
    target_token = target_token[target_pos].view(-1)
    loss         = torch.nn.functional.cross_entropy(pred_logits, target_token)

    engine.backward(loss)
    engine.step()

Expected behavior Fix the error Trying to backward through the graph a second time

ds_report output

[2023-10-17 22:28:50,665] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch']
torch version .................... 2.1.0+cu121
deepspeed install path ........... ['/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.11.1, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7
shared memory (/dev/shm) size .... 125.77 GB

Screenshots No Screenshots

System info (please complete the following information): OS: Ubuntu 20.04.4 GPU: 1x RTX 3090ti Interconnects: No Python version: 3.11.4

Launcher context PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 deepspeed --include localhost:0 train.py

Docker context No docker

Additional context Complete deepspeed output In Reproduce, I simplified some of the code, so the error traceback is slightly different from Reproduce.

[2023-10-17 22:38:23,225] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:24,055] [WARNING] [runner.py:203:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2023-10-17 22:38:24,056] [INFO] [runner.py:570:main] cmd = /home/user/miniconda3/envs/llama2/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbNV19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None train.py
[2023-10-17 22:38:26,035] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:26,860] [INFO] [launch.py:145:main] WORLD INFO DICT: {'localhost': [5]}
[2023-10-17 22:38:26,860] [INFO] [launch.py:151:main] nnodes=1, num_local_procs=1, node_rank=0
[2023-10-17 22:38:26,860] [INFO] [launch.py:162:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0]})
[2023-10-17 22:38:26,860] [INFO] [launch.py:163:main] dist_world_size=1
[2023-10-17 22:38:26,860] [INFO] [launch.py:165:main] Setting CUDA_VISIBLE_DEVICES=5
[2023-10-17 22:38:28,800] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:31,345] [INFO] [comm.py:637:init_distributed] cdb=None
[2023-10-17 22:38:31,346] [INFO] [comm.py:668:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2023-10-17 22:41:14,468] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.11.1, git-hash=unknown, git-branch=unknown
[2023-10-17 22:41:20,391] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
 [WARNING]  cpu_adam cuda is missing or is incompatible with installed torch, only cpu ops can be compiled!
Using /home/user/.cache/torch_extensions/py311_cu121 as PyTorch extensions root...
Emitting ninja build file /home/user/.cache/torch_extensions/py311_cu121/cpu_adam/build.ninja...
Building extension module cpu_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module cpu_adam...
Time to load cpu_adam op: 2.476133108139038 seconds
Adam Optimizer #0 is created with AVX512 arithmetic capability.
Config: alpha=0.000020, betas=(0.900000, 0.950000), weight_decay=0.000100, adam_w=1
[2023-10-17 22:41:25,279] [INFO] [logging.py:96:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer
[2023-10-17 22:41:25,280] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = DeepSpeedCPUAdam
[2023-10-17 22:41:25,313] [INFO] [utils.py:56:is_zero_supported_optimizer] Checking ZeRO support for optimizer=DeepSpeedCPUAdam type=<class 'deepspeed.ops.adam.cpu_adam.DeepSpeedCPUAdam'>
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] Creating fp16 ZeRO stage 3 optimizer, MiCS is enabled False, Hierarchical params gather False
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 3 optimizer
[2023-10-17 22:41:25,766] [INFO] [utils.py:802:see_memory_usage] Stage 3 initialize beginning
[2023-10-17 22:41:25,767] [INFO] [utils.py:803:see_memory_usage] MA 12.99 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:25,767] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 186.86 GB, percent = 74.3%
[2023-10-17 22:41:25,771] [INFO] [stage3.py:126:__init__] Reduce bucket size 50000000
[2023-10-17 22:41:25,771] [INFO] [stage3.py:127:__init__] Prefetch bucket size 50,000,000
[2023-10-17 22:41:26,225] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [begin]
[2023-10-17 22:41:26,225] [INFO] [utils.py:803:see_memory_usage] MA 12.99 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:26,225] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 187.52 GB, percent = 74.6%
Parameter Offload: Total persistent parameters: 10247936 in 260 params
[2023-10-17 22:41:37,268] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [end]
[2023-10-17 22:41:37,268] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:37,269] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 199.51 GB, percent = 79.3%
[2023-10-17 22:41:37,759] [INFO] [utils.py:802:see_memory_usage] Before creating fp16 partitions
[2023-10-17 22:41:37,760] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:37,760] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.1 GB, percent = 79.6%
[2023-10-17 22:41:38,377] [INFO] [utils.py:802:see_memory_usage] After creating fp16 partitions: 1
[2023-10-17 22:41:38,377] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:38,377] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 199.63 GB, percent = 79.4%
[2023-10-17 22:41:38,946] [INFO] [utils.py:802:see_memory_usage] Before creating fp32 partitions
[2023-10-17 22:41:38,947] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:38,947] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.27 GB, percent = 79.6%
[2023-10-17 22:41:39,533] [INFO] [utils.py:802:see_memory_usage] After creating fp32 partitions
[2023-10-17 22:41:39,533] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:39,533] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.19 GB, percent = 79.6%
[2023-10-17 22:41:40,012] [INFO] [utils.py:802:see_memory_usage] Before initializing optimizer states
[2023-10-17 22:41:40,012] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:40,013] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.68 GB, percent = 79.8%
[2023-10-17 22:41:40,636] [INFO] [utils.py:802:see_memory_usage] After initializing optimizer states
[2023-10-17 22:41:40,637] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:40,653] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.93 GB, percent = 79.9%
[2023-10-17 22:41:40,653] [INFO] [stage3.py:459:_setup_for_real_optimizer] optimizer state initialized
[2023-10-17 22:41:41,334] [INFO] [utils.py:802:see_memory_usage] After initializing ZeRO optimizer
[2023-10-17 22:41:41,335] [INFO] [utils.py:803:see_memory_usage] MA 0.41 GB         Max_MA 0.41 GB         CA 13.1 GB         Max_CA 13 GB 
[2023-10-17 22:41:41,335] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 201.63 GB, percent = 80.2%
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = adamw
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = WarmupDecayLR
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupDecayLR object at 0x7fb1e7f0ed50>
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[2e-05], mom=[[0.9, 0.95]]
[2023-10-17 22:41:41,337] [INFO] [config.py:968:print] DeepSpeedEngine configuration:
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   activation_checkpointing_config  {
    "partition_activations": true, 
    "contiguous_memory_optimization": true, 
    "cpu_checkpointing": true, 
    "number_checkpoints": null, 
    "synchronize_checkpoint_boundary": false, 
    "profile": false
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   amp_enabled .................. False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   amp_params ................... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   autotuning_config ............ {
    "enabled": false, 
    "start_step": null, 
    "end_step": null, 
    "metric_path": null, 
    "arg_mappings": null, 
    "metric": "throughput", 
    "model_info": null, 
    "results_dir": "autotuning_results", 
    "exps_dir": "autotuning_exps", 
    "overwrite": true, 
    "fast": true, 
    "start_profile_step": 3, 
    "end_profile_step": 5, 
    "tuner_type": "gridsearch", 
    "tuner_early_stopping": 5, 
    "tuner_num_trials": 50, 
    "model_info_path": null, 
    "mp_size": 1, 
    "max_train_batch_size": null, 
    "min_train_batch_size": 1, 
    "max_train_micro_batch_size_per_gpu": 1.024000e+03, 
    "min_train_micro_batch_size_per_gpu": 1, 
    "num_tuning_micro_batch_sizes": 3
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   bfloat16_enabled ............. True
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_parallel_write_pipeline  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_tag_validation_enabled  True
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_tag_validation_fail  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7fb547229a90>
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   communication_data_type ...... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   curriculum_enabled_legacy .... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   curriculum_params_legacy ..... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   data_efficiency_enabled ...... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dataloader_drop_last ......... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   disable_allgather ............ False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dump_state ................... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dynamic_loss_scale_args ...... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_enabled ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_gas_boundary_resolution  1
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_layer_name ........ bert.encoder.layer
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_layer_num ......... 0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_max_iter .......... 100
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_stability ......... 1e-06
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_tol ............... 0.01
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_verbose ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   elasticity_enabled ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   flops_profiler_config ........ {
    "enabled": false, 
    "recompute_fwd_factor": 0.0, 
    "profile_step": 1, 
    "module_depth": -1, 
    "top_modules": 1, 
    "detailed": true, 
    "output_file": null
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_auto_cast ............... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_enabled ................. False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_master_weights_and_gradients  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   global_rank .................. 0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   grad_accum_dtype ............. bf16
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_accumulation_steps .. 8
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_clipping ............ 1.0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_predivide_factor .... 1.0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   initial_dynamic_scale ........ 1
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   load_universal_checkpoint .... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   loss_scale ................... 1.0
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   memory_breakdown ............. False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   mics_hierarchial_params_gather  False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   mics_shard_size .............. -1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') enabled=False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   nebula_config ................ {
    "enabled": false, 
    "persistent_storage_path": null, 
    "persistent_time_interval": 100, 
    "num_of_version_in_retention": 2, 
    "enable_nebula_load": true, 
    "load_path": null
}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_legacy_fusion ...... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_name ............... adamw
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_params ............. {'lr': 2e-05, 'betas': [0.9, 0.95], 'eps': 1e-08, 'weight_decay': 0.0001}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pld_enabled .................. False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pld_params ................... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   prescale_gradients ........... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   scheduler_name ............... WarmupDecayLR
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 2e-05, 'warmup_num_steps': 1000, 'total_num_steps': 10000}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   sparse_attention ............. None
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   sparse_gradients_enabled ..... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   steps_per_print .............. 10
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   train_batch_size ............. 8
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   train_micro_batch_size_per_gpu  1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   use_node_local_storage ....... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   wall_clock_breakdown ......... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   weight_quantization_config ... None
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   world_size ................... 1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_allow_untested_optimizer  False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_config .................. stage=3 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=50000000 allgather_partitions=True allgather_bucket_size=50000000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=DeepSpeedZeroOffloadParamConfig(device='cpu', nvme_path=None, buffer_count=5, buffer_size=100,000,000, max_in_cpu=1,000,000,000, pin_memory=False) offload_optimizer=DeepSpeedZeroOffloadOptimizerConfig(device='cpu', nvme_path=None, buffer_count=4, pin_memory=False, pipeline=False, pipeline_read=False, pipeline_write=False, fast_init=False) sub_group_size=1,000,000,000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50,000,000 param_persistence_threshold=100,000 model_persistence_threshold=sys.maxsize max_live_parameters=1,000,000,000 max_reuse_distance=1,000,000,000 gather_16bit_weights_on_model_save=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_enabled ................. True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_force_ds_cpu_optimizer .. True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_optimization_stage ...... 3
[2023-10-17 22:41:41,339] [INFO] [config.py:958:print_user_config]   json = {
    "train_micro_batch_size_per_gpu": 1, 
    "gradient_accumulation_steps": 8, 
    "gradient_clipping": 1.0, 
    "steps_per_print": 10, 
    "optimizer": {
        "type": "AdamW", 
        "params": {
            "lr": 2e-05, 
            "betas": [0.9, 0.95], 
            "eps": 1e-08, 
            "weight_decay": 0.0001
        }
    }, 
    "scheduler": {
        "type": "WarmupDecayLR", 
        "params": {
            "warmup_min_lr": 0, 
            "warmup_max_lr": 2e-05, 
            "warmup_num_steps": 1000, 
            "total_num_steps": 1.000000e+04
        }
    }, 
    "bf16": {
        "enabled": true
    }, 
    "zero_optimization": {
        "stage": 3, 
        "offload_param": {
            "device": "cpu", 
            "pin_memory": false
        }, 
        "offload_optimizer": {
            "device": "cpu", 
            "pin_memory": false
        }, 
        "allgather_bucket_size": 5.000000e+07, 
        "reduce_bucket_size": 5.000000e+07, 
        "allgather_partitions": true, 
        "reduce_scatter": true, 
        "overlap_comm": false, 
        "contiguous_gradients": true
    }, 
    "activation_checkpointing": {
        "partition_activations": true, 
        "cpu_checkpointing": true, 
        "contiguous_memory_optimization": true, 
        "number_checkpoints": null, 
        "synchronize_checkpoint_boundary": false, 
        "profile": false
    }, 
    "data_types": {
        "grad_accum_dtype": "bf16"
    }
}

Traceback (most recent call last):
  File "/data/user/llama2/src/7_train/train.py", line 449, in <module>
    trainer.train_train()
  File "/data/user/llama2/src/7_train/train.py", line 408, in train_train
    self.train_epoch(epoch)
  File "/data/user/llama2/src/7_train/train.py", line 437, in train_epoch
    self.model_engine.backward(loss)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1929, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2091, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
[2023-10-17 22:42:02,884] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 230633
[2023-10-17 22:42:02,885] [ERROR] [launch.py:321:sigkill_handler] ['/home/user/miniconda3/envs/llama2/bin/python', '-u', 'train.py', '--local_rank=0'] exits with return code = 1

I would appreciate any help.

MIkumikumi0116 commented 11 months ago

Describe the bug When I am training my model, everything goes well in the first batch, but an error occurs in the second batch: RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.

To Reproduce I am fine-tuning llama2 using my own implementation of LoRA. Each lora linear in my llama2 has different LoRA rank. Each linear in my model is replaced to

class Adaptive_Linear(torch.nn.Module):
    def __init__(
        self,
        model_arg: dict, # batch size, max seqlen, vocab size, etc.
        name: str, # example: layer_list.0.attention.key
        lora_rank_dict: dict, # lora rank of each linear
        in_features: int,
        out_features: int 
    ):
        super().__init__()

        self.name      = name
        self.model_arg = model_arg

        self.linear    = torch.nn.Linear(in_features, out_features, bias = False)

        self.lora_rank = lora_rank_dict[self.name]
        if self.lora_rank > 0:
            self.lora_dropout  = torch.nn.Dropout(model_arg.lora_dropout)
            self.lora_a_linear = torch.nn.Linear(in_features, self.lora_rank,  bias = False)
            self.lora_b_linear = torch.nn.Linear(self.lora_rank, out_features, bias = False)

            self.lora_a_linear.weight.data = torch.rand_like(self.lora_a_linear.weight.data) - 1 * 2
            self.lora_b_linear.weight.data = torch.zeros_like(self.lora_b_linear.weight.data)

    def trainable_parameters(self):
        # only train lora matrix
        # the trainable parameters of my model are trainable_parameters
        # of all Adaptive_Linear connected together using itertools.chain.
        return iter((self.lora_a_linear.weight, self.lora_b_linear.weight))

    def forward(self, x):
        hidden = self.linear(x)

        if self.lora_rank > 0:
            input  = self.lora_dropout(x)
            lora_a = self.lora_a_linear(input)
            lora_b = self.lora_b_linear(lora_a)
            output = hidden + lora_b
        else:
            output = hidden

        return output

My dataset is

class Dataset(torch.utils.data.Dataset):
  def __init__(self, model_arg, dataset_path):
      self.model_arg    = model_arg

      self.dataset_path = dataset_path
      self.dataset      = torch.load(self.dataset_path) # List[Tuple] list of tuples of (input token list, output token list, label) token and label are int

      self.dataset      = [data_item for data_item in self.dataset if len(data_item[0]) >= MIN_SEQ_LEN and len(data_item[0]) <= MAX_SEQ_LEN]

  def __len__(self):
      return len(self.dataset)

  def __getitem__(self, index):
      input_token, output_token, label = self.dataset

      return input_token, output_token, label

Collate function to pad tokens of a batch to same length

collate_fn(batch):
    input_token_list, output_token_list, label_list = zip(*batch)

    padded_input_token  = torch.nn.utils.rnn.pad_sequence([seq for seq in input_token_list],  batch_first = True)
    padded_output_token = torch.nn.utils.rnn.pad_sequence([seq for seq in output_token_list], batch_first = True)

    return padded_input_token, padded_output_token, label_list

Initialize model and deepspeed

model = Llama_2_With_Adaptive_Lora(model_arg, lora_rank_dict)
model.load_state_dict(state_dict)
for name, parameter in model.named_parameters():
    if not ('lora' in name):
        parameter.requires_grad = False

engine, optimizer, train_dataloader, _= deepspeed.initialize(
    model            = model,
    model_parameters = model.trainable_parameters(),
    training_data    = train_dataset,
    args             = args, # local_rank only
    collate_fn       = collate_fn,
    config           = deepspeed_config
)

Deepspeed config

{
    "train_micro_batch_size_per_gpu": 1,
    "gradient_accumulation_steps": 8,
    "gradient_clipping": 1.0,
    "steps_per_print": 10,
    "wall_clock_breakdown": false,
    "memory_breakdown": false,
    "optimizer": {
        "type": "AdamW",
        "params": {
            "lr": 2e-5,
            "betas": [
                0.9,
                0.95
            ],
            "eps": 1e-8,
            "weight_decay": 1e-4
        }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": {
            "warmup_min_lr": 0,
            "warmup_max_lr": 2e-5,
            "warmup_num_steps": 1000,
            "total_num_steps": 10000
        }
    },
    "bf16": {
        "enabled": true,
        "loss_scale": 0,
        "loss_scale_window": 1000,
        "initial_scale_power": 16,
        "hysteresis": 2,
        "min_loss_scale": 1
    },
    "zero_optimization": {
        "stage": 3,
        "offload_param": {
            "device": "cpu",
            "pin_memory": true
        },
        "offload_optimizer": {
            "device": "cpu",
            "pin_memory": true
        },

        "allgather_partitions": true,
        "allgather_bucket_size": 1e7,
        "reduce_scatter": true,
        "reduce_bucket_size": 1e7,

        "overlap_comm": true,
        "contiguous_gradients": true,

        "sub_group_size": 1e9,
        "stage3_prefetch_bucket_size": 1e7,
        "stage3_param_persistence_threshold": 1e5,
        "stage3_max_live_parameters": 1e8,
        "stage3_max_reuse_distance": 1e8,
        "stage3_gather_16bit_weights_on_model_save": true
    },
    "activation_checkpointing": {
        "partition_activations": true,
        "number_checkpoints": 100,

        "cpu_checkpointing": true,
        "contiguous_memory_optimization": true,
        "synchronize_checkpoint_boundary": false,
        "profile": true
    },
    "data_types": {
        "grad_accum_dtype":"bf16"
    }
}

Main code of training

engine.train()
for input_token, target_token, _ in train_dataloader:
    input_token  = input_token. to(engine.device)
    target_token = target_token.to(engine.device)

    pred_logits  = engine(input_token)

    target_pos   = (target_token != PAD_ID) # only calculate loss on non-pad token
    pred_logits  = pred_logits[ target_pos].view(-1, model_arg.vocab_size)
    target_token = target_token[target_pos].view(-1)
    loss         = torch.nn.functional.cross_entropy(pred_logits, target_token)

    engine.backward(loss)
    engine.step()

Expected behavior Fix the error Trying to backward through the graph a second time

ds_report output

[2023-10-17 22:28:50,665] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
 [WARNING]  Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.1
 [WARNING]  using untested triton version (2.1.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch']
torch version .................... 2.1.0+cu121
deepspeed install path ........... ['/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed']
deepspeed info ................... 0.11.1, unknown, unknown
torch cuda version ............... 12.1
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7
shared memory (/dev/shm) size .... 125.77 GB

Screenshots No Screenshots

System info (please complete the following information): OS: Ubuntu 20.04.4 GPU: 1x RTX 3090ti Interconnects: No Python version: 3.11.4

Launcher context PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:64 deepspeed --include localhost:0 train.py

Docker context No docker

Additional context Complete deepspeed output In Reproduce, I simplified some of the code, so the error traceback is slightly different from Reproduce.

[2023-10-17 22:38:23,225] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:24,055] [WARNING] [runner.py:203:fetch_hostfile] Unable to find hostfile, will proceed with training with local resources only.
[2023-10-17 22:38:24,056] [INFO] [runner.py:570:main] cmd = /home/user/miniconda3/envs/llama2/bin/python -u -m deepspeed.launcher.launch --world_info=eyJsb2NhbGhvc3QiOiBbNV19 --master_addr=127.0.0.1 --master_port=29500 --enable_each_rank_log=None train.py
[2023-10-17 22:38:26,035] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:26,860] [INFO] [launch.py:145:main] WORLD INFO DICT: {'localhost': [5]}
[2023-10-17 22:38:26,860] [INFO] [launch.py:151:main] nnodes=1, num_local_procs=1, node_rank=0
[2023-10-17 22:38:26,860] [INFO] [launch.py:162:main] global_rank_mapping=defaultdict(<class 'list'>, {'localhost': [0]})
[2023-10-17 22:38:26,860] [INFO] [launch.py:163:main] dist_world_size=1
[2023-10-17 22:38:26,860] [INFO] [launch.py:165:main] Setting CUDA_VISIBLE_DEVICES=5
[2023-10-17 22:38:28,800] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2023-10-17 22:38:31,345] [INFO] [comm.py:637:init_distributed] cdb=None
[2023-10-17 22:38:31,346] [INFO] [comm.py:668:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2023-10-17 22:41:14,468] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.11.1, git-hash=unknown, git-branch=unknown
[2023-10-17 22:41:20,391] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
�[93m [WARNING] �[0m cpu_adam cuda is missing or is incompatible with installed torch, only cpu ops can be compiled!
Using /home/user/.cache/torch_extensions/py311_cu121 as PyTorch extensions root...
Emitting ninja build file /home/user/.cache/torch_extensions/py311_cu121/cpu_adam/build.ninja...
Building extension module cpu_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
ninja: no work to do.
Loading extension module cpu_adam...
Time to load cpu_adam op: 2.476133108139038 seconds
Adam Optimizer #0 is created with AVX512 arithmetic capability.
Config: alpha=0.000020, betas=(0.900000, 0.950000), weight_decay=0.000100, adam_w=1
[2023-10-17 22:41:25,279] [INFO] [logging.py:96:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer
[2023-10-17 22:41:25,280] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = DeepSpeedCPUAdam
[2023-10-17 22:41:25,313] [INFO] [utils.py:56:is_zero_supported_optimizer] Checking ZeRO support for optimizer=DeepSpeedCPUAdam type=<class 'deepspeed.ops.adam.cpu_adam.DeepSpeedCPUAdam'>
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] Creating fp16 ZeRO stage 3 optimizer, MiCS is enabled False, Hierarchical params gather False
[2023-10-17 22:41:25,313] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 3 optimizer
[2023-10-17 22:41:25,766] [INFO] [utils.py:802:see_memory_usage] Stage 3 initialize beginning
[2023-10-17 22:41:25,767] [INFO] [utils.py:803:see_memory_usage] MA 12.99 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:25,767] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 186.86 GB, percent = 74.3%
[2023-10-17 22:41:25,771] [INFO] [stage3.py:126:__init__] Reduce bucket size 50000000
[2023-10-17 22:41:25,771] [INFO] [stage3.py:127:__init__] Prefetch bucket size 50,000,000
[2023-10-17 22:41:26,225] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [begin]
[2023-10-17 22:41:26,225] [INFO] [utils.py:803:see_memory_usage] MA 12.99 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:26,225] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 187.52 GB, percent = 74.6%
Parameter Offload: Total persistent parameters: 10247936 in 260 params
[2023-10-17 22:41:37,268] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [end]
[2023-10-17 22:41:37,268] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 12.99 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:37,269] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 199.51 GB, percent = 79.3%
[2023-10-17 22:41:37,759] [INFO] [utils.py:802:see_memory_usage] Before creating fp16 partitions
[2023-10-17 22:41:37,760] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:37,760] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.1 GB, percent = 79.6%
[2023-10-17 22:41:38,377] [INFO] [utils.py:802:see_memory_usage] After creating fp16 partitions: 1
[2023-10-17 22:41:38,377] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:38,377] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 199.63 GB, percent = 79.4%
[2023-10-17 22:41:38,946] [INFO] [utils.py:802:see_memory_usage] Before creating fp32 partitions
[2023-10-17 22:41:38,947] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:38,947] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.27 GB, percent = 79.6%
[2023-10-17 22:41:39,533] [INFO] [utils.py:802:see_memory_usage] After creating fp32 partitions
[2023-10-17 22:41:39,533] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:39,533] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.19 GB, percent = 79.6%
[2023-10-17 22:41:40,012] [INFO] [utils.py:802:see_memory_usage] Before initializing optimizer states
[2023-10-17 22:41:40,012] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:40,013] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.68 GB, percent = 79.8%
[2023-10-17 22:41:40,636] [INFO] [utils.py:802:see_memory_usage] After initializing optimizer states
[2023-10-17 22:41:40,637] [INFO] [utils.py:803:see_memory_usage] MA 0.31 GB         Max_MA 0.31 GB         CA 13.01 GB         Max_CA 13 GB 
[2023-10-17 22:41:40,653] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 200.93 GB, percent = 79.9%
[2023-10-17 22:41:40,653] [INFO] [stage3.py:459:_setup_for_real_optimizer] optimizer state initialized
[2023-10-17 22:41:41,334] [INFO] [utils.py:802:see_memory_usage] After initializing ZeRO optimizer
[2023-10-17 22:41:41,335] [INFO] [utils.py:803:see_memory_usage] MA 0.41 GB         Max_MA 0.41 GB         CA 13.1 GB         Max_CA 13 GB 
[2023-10-17 22:41:41,335] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 201.63 GB, percent = 80.2%
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = adamw
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = WarmupDecayLR
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupDecayLR object at 0x7fb1e7f0ed50>
[2023-10-17 22:41:41,335] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[2e-05], mom=[[0.9, 0.95]]
[2023-10-17 22:41:41,337] [INFO] [config.py:968:print] DeepSpeedEngine configuration:
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   activation_checkpointing_config  {
    "partition_activations": true, 
    "contiguous_memory_optimization": true, 
    "cpu_checkpointing": true, 
    "number_checkpoints": null, 
    "synchronize_checkpoint_boundary": false, 
    "profile": false
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   amp_enabled .................. False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   amp_params ................... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   autotuning_config ............ {
    "enabled": false, 
    "start_step": null, 
    "end_step": null, 
    "metric_path": null, 
    "arg_mappings": null, 
    "metric": "throughput", 
    "model_info": null, 
    "results_dir": "autotuning_results", 
    "exps_dir": "autotuning_exps", 
    "overwrite": true, 
    "fast": true, 
    "start_profile_step": 3, 
    "end_profile_step": 5, 
    "tuner_type": "gridsearch", 
    "tuner_early_stopping": 5, 
    "tuner_num_trials": 50, 
    "model_info_path": null, 
    "mp_size": 1, 
    "max_train_batch_size": null, 
    "min_train_batch_size": 1, 
    "max_train_micro_batch_size_per_gpu": 1.024000e+03, 
    "min_train_micro_batch_size_per_gpu": 1, 
    "num_tuning_micro_batch_sizes": 3
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   bfloat16_enabled ............. True
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_parallel_write_pipeline  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_tag_validation_enabled  True
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   checkpoint_tag_validation_fail  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7fb547229a90>
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   communication_data_type ...... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   curriculum_enabled_legacy .... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   curriculum_params_legacy ..... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   data_efficiency_enabled ...... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dataloader_drop_last ......... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   disable_allgather ............ False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dump_state ................... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   dynamic_loss_scale_args ...... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_enabled ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_gas_boundary_resolution  1
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_layer_name ........ bert.encoder.layer
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_layer_num ......... 0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_max_iter .......... 100
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_stability ......... 1e-06
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_tol ............... 0.01
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   eigenvalue_verbose ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   elasticity_enabled ........... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   flops_profiler_config ........ {
    "enabled": false, 
    "recompute_fwd_factor": 0.0, 
    "profile_step": 1, 
    "module_depth": -1, 
    "top_modules": 1, 
    "detailed": true, 
    "output_file": null
}
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_auto_cast ............... None
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_enabled ................. False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   fp16_master_weights_and_gradients  False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   global_rank .................. 0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   grad_accum_dtype ............. bf16
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_accumulation_steps .. 8
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_clipping ............ 1.0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   gradient_predivide_factor .... 1.0
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   initial_dynamic_scale ........ 1
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   load_universal_checkpoint .... False
[2023-10-17 22:41:41,338] [INFO] [config.py:972:print]   loss_scale ................... 1.0
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   memory_breakdown ............. False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   mics_hierarchial_params_gather  False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   mics_shard_size .............. -1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') enabled=False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   nebula_config ................ {
    "enabled": false, 
    "persistent_storage_path": null, 
    "persistent_time_interval": 100, 
    "num_of_version_in_retention": 2, 
    "enable_nebula_load": true, 
    "load_path": null
}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_legacy_fusion ...... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_name ............... adamw
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   optimizer_params ............. {'lr': 2e-05, 'betas': [0.9, 0.95], 'eps': 1e-08, 'weight_decay': 0.0001}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pld_enabled .................. False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   pld_params ................... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   prescale_gradients ........... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   scheduler_name ............... WarmupDecayLR
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 2e-05, 'warmup_num_steps': 1000, 'total_num_steps': 10000}
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   sparse_attention ............. None
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   sparse_gradients_enabled ..... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   steps_per_print .............. 10
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   train_batch_size ............. 8
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   train_micro_batch_size_per_gpu  1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   use_node_local_storage ....... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   wall_clock_breakdown ......... False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   weight_quantization_config ... None
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   world_size ................... 1
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_allow_untested_optimizer  False
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_config .................. stage=3 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=50000000 allgather_partitions=True allgather_bucket_size=50000000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=DeepSpeedZeroOffloadParamConfig(device='cpu', nvme_path=None, buffer_count=5, buffer_size=100,000,000, max_in_cpu=1,000,000,000, pin_memory=False) offload_optimizer=DeepSpeedZeroOffloadOptimizerConfig(device='cpu', nvme_path=None, buffer_count=4, pin_memory=False, pipeline=False, pipeline_read=False, pipeline_write=False, fast_init=False) sub_group_size=1,000,000,000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50,000,000 param_persistence_threshold=100,000 model_persistence_threshold=sys.maxsize max_live_parameters=1,000,000,000 max_reuse_distance=1,000,000,000 gather_16bit_weights_on_model_save=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_enabled ................. True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_force_ds_cpu_optimizer .. True
[2023-10-17 22:41:41,339] [INFO] [config.py:972:print]   zero_optimization_stage ...... 3
[2023-10-17 22:41:41,339] [INFO] [config.py:958:print_user_config]   json = {
    "train_micro_batch_size_per_gpu": 1, 
    "gradient_accumulation_steps": 8, 
    "gradient_clipping": 1.0, 
    "steps_per_print": 10, 
    "optimizer": {
        "type": "AdamW", 
        "params": {
            "lr": 2e-05, 
            "betas": [0.9, 0.95], 
            "eps": 1e-08, 
            "weight_decay": 0.0001
        }
    }, 
    "scheduler": {
        "type": "WarmupDecayLR", 
        "params": {
            "warmup_min_lr": 0, 
            "warmup_max_lr": 2e-05, 
            "warmup_num_steps": 1000, 
            "total_num_steps": 1.000000e+04
        }
    }, 
    "bf16": {
        "enabled": true
    }, 
    "zero_optimization": {
        "stage": 3, 
        "offload_param": {
            "device": "cpu", 
            "pin_memory": false
        }, 
        "offload_optimizer": {
            "device": "cpu", 
            "pin_memory": false
        }, 
        "allgather_bucket_size": 5.000000e+07, 
        "reduce_bucket_size": 5.000000e+07, 
        "allgather_partitions": true, 
        "reduce_scatter": true, 
        "overlap_comm": false, 
        "contiguous_gradients": true
    }, 
    "activation_checkpointing": {
        "partition_activations": true, 
        "cpu_checkpointing": true, 
        "contiguous_memory_optimization": true, 
        "number_checkpoints": null, 
        "synchronize_checkpoint_boundary": false, 
        "profile": false
    }, 
    "data_types": {
        "grad_accum_dtype": "bf16"
    }
}

Traceback (most recent call last):
  File "/data/user/llama2/src/7_train/train.py", line 449, in <module>
    trainer.train_train()
  File "/data/user/llama2/src/7_train/train.py", line 408, in train_train
    self.train_epoch(epoch)
  File "/data/user/llama2/src/7_train/train.py", line 437, in train_epoch
    self.model_engine.backward(loss)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 1929, in backward
    self.optimizer.backward(loss, retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
    ret_val = func(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/zero/stage3.py", line 2091, in backward
    self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
    scaled_loss.backward(retain_graph=retain_graph)
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch/_tensor.py", line 492, in backward
    torch.autograd.backward(
  File "/home/user/miniconda3/envs/llama2/lib/python3.11/site-packages/torch/autograd/__init__.py", line 251, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: Trying to backward through the graph a second time (or directly access saved tensors after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved tensors after calling backward.
[2023-10-17 22:42:02,884] [INFO] [launch.py:315:sigkill_handler] Killing subprocess 230633
[2023-10-17 22:42:02,885] [ERROR] [launch.py:321:sigkill_handler] ['/home/user/miniconda3/envs/llama2/bin/python', '-u', 'train.py', '--local_rank=0'] exits with return code = 1

I would appreciate any help.

I should submit the issue to https://github.com/microsoft/DeepSpeed Apologies for the error submission.