Error while finetune with deepspeed zero3 for multi-node (multi-GPUS)

[hahmad2008]

What is the accelerate.yaml and the config.yaml configuration to finetune on multi-nodes and each node with multi gpus?
I used the following configs but I get this error:

• axolotl commit: git checkout -f 268543a3be20aa8de662761bb36645cbb5ff8f5b
• command:

 accelerate launch  --config_file accelerate.yaml  -m  axolotl.cli.train config.yaml  --deepspeed deepspeed_configs/zero3_bf16.json 

--
NOTE: the following was working with this old axolotl commit: git checkout -f 4de4b4089fbc43e04c696c10c1ac51eebf2c6a99

• accelerate.yaml

compute_environment: LOCAL_MACHINE
debug: false
distributed_type: 'MULTI_GPU'
downcast_bf16: false
enable_cpu_affinity: false
machine_rank: 0
main_process_ip: NODE-IP
main_process_port: 8000
main_training_function: main
mixed_precision: 'bf16'
num_machines: 10
num_processes: 80
rdzv_backend: static
same_network: true
tpu_use_cluster: false
tpu_use_sudo: false
use_cpu: false

• machine_rank: 0 for master node and other 3 nodes with rank 1,2,3. and we run the same command but with diff accelerate.yaml for each node.
• they have common config.yaml

• config.yaml

base_model: meta-llama/Llama-3.1-70B-Instruct
base_model_config: meta-llama/Llama-3.1-70B-Instruct
model_type: LlamaForCausalLM
tokenizer_type: AutoTokenizer
load_in_8bit: false
load_in_4bit: false
strict: false
chat_template: llama3
datasets:
  - path: dataset.jsonl
    type: chat_template
    conversation: llama3
    field_messages: conversations
    message_field_role: from
    message_field_content: value
    roles:
      system:
        - system
      user:
        - human
      assistant:
        - gpt
dataset_prepared_path: /shared_ds/prep-data
output_dir: /shared_ds/out-dir
val_set_size: 0.01
sequence_len: 8192
sample_packing: true
eval_sample_packing: false
pad_to_sequence_len: true
wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:
hub_model_id:
gradient_accumulation_steps: 1
micro_batch_size: 1
num_epochs: 2
optimizer: paged_adamw_32bit
lr_scheduler: cosine
learning_rate: 2e-4
train_on_inputs: false
group_by_length: false
bf16: auto
fp16:
tf32: false
gradient_checkpointing: true
gradient_checkpointing_kwargs:
  use_reentrant: false
early_stopping_patience:
resume_from_checkpoint:
logging_steps: 1
xformers_attention:
flash_attention: true
adapter: lora
lora_model_dir:
lora_r: 32
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true
lora_fan_in_fan_out:
warmup_ratio: 0.02
eval_table_size:
eval_steps: 1000
saves_per_epoch: 5
debug:
deepspeed: deepspeed_configs/zero3.json
weight_decay: 0.0
fsdp:
fsdp_config:
special_tokens:
  pad_token: <|end_of_text|>

• deepspeed_configs/zero3.json

{
  "zero_optimization": {
    "stage": 3,
    "overlap_comm": true,
    "contiguous_gradients": true,
    "sub_group_size": 0,
    "reduce_bucket_size": "auto",
    "stage3_prefetch_bucket_size": "auto",
    "stage3_param_persistence_threshold": "auto",
    "stage3_max_live_parameters": 0,
    "stage3_max_reuse_distance": 0,
    "stage3_gather_16bit_weights_on_model_save": true
  },
  "bf16": {
    "enabled": true
  },
  "gradient_accumulation_steps": "auto",
  "gradient_clipping": "auto",
  "train_batch_size": "auto",
  "train_micro_batch_size_per_gpu": "auto",
  "wall_clock_breakdown": false
}

• Error:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
[rank0]:     return _run_code(code, main_globals, None,
[rank0]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
[rank0]:     exec(code, run_globals)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 113, in <module>
[rank0]:     fire.Fire(do_cli)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 135, in Fire
[rank0]:     component_trace = _Fire(component, args, parsed_flag_args, context, name)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 468, in _Fire
[rank0]:     component, remaining_args = _CallAndUpdateTrace(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 684, in _CallAndUpdateTrace
[rank0]:     component = fn(*varargs, **kwargs)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 87, in do_cli
[rank0]:     return do_train(parsed_cfg, parsed_cli_args)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 44, in do_train
[rank0]:     model, tokenizer = train(cfg=cfg, dataset_meta=dataset_meta)
[rank0]:   File "axolotl/src/axolotl/train.py", line 187, in train
[rank0]:     trainer.train(resume_from_checkpoint=resume_from_checkpoint)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2171, in train
[rank0]:     return inner_training_loop(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2531, in _inner_training_loop
[rank0]: Traceback (most recent call last):
[rank0]:   File "/usr/lib/python3.10/runpy.py", line 196, in _run_module_as_main
[rank0]:     return _run_code(code, main_globals, None,
[rank0]:   File "/usr/lib/python3.10/runpy.py", line 86, in _run_code
[rank0]:     exec(code, run_globals)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 113, in <module>
[rank0]:     fire.Fire(do_cli)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 135, in Fire
[rank0]:     component_trace = _Fire(component, args, parsed_flag_args, context, name)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 468, in _Fire
[rank0]:     component, remaining_args = _CallAndUpdateTrace(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/fire/core.py", line 684, in _CallAndUpdateTrace
[rank0]:     component = fn(*varargs, **kwargs)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 87, in do_cli
[rank0]:     return do_train(parsed_cfg, parsed_cli_args)
[rank0]:   File "axolotl/src/axolotl/cli/train.py", line 44, in do_train
[rank0]:     model, tokenizer = train(cfg=cfg, dataset_meta=dataset_meta)
[rank0]:   File "axolotl/src/axolotl/train.py", line 187, in train
[rank0]:     trainer.train(resume_from_checkpoint=resume_from_checkpoint)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2171, in train
[rank0]:     return inner_training_loop(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 2531, in _inner_training_loop
[rank0]:     self.loss_scaler.backward(loss.float(), retain_graph=retain_graph)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/fp16/loss_scaler.py", line 63, in backward
[rank0]:     scaled_loss.backward(retain_graph=retain_graph)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/_tensor.py", line 581, in backward
[rank0]:     torch.autograd.backward(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/__init__.py", line 347, in backward
[rank0]:     _engine_run_backward(
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
[rank0]:     return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/autograd/function.py", line 307, in apply
[rank0]:     return user_fn(self, *args)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/amp/autocast_mode.py", line 511, in decorate_bwd
[rank0]:     return bwd(*args, **kwargs)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/deepspeed/runtime/zero/linear.py", line 80, in backward
[rank0]:     input, weight, bias = ctx.saved_tensors
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py", line 1129, in unpack_hook
[rank0]:     frame.check_recomputed_tensors_match(gid)
[rank0]:   File "/usr/local/lib/python3.10/dist-packages/torch/utils/checkpoint.py", line 903, in check_recomputed_tensors_match
[rank0]:     raise CheckpointError(
[rank0]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
[rank0]: tensor at position 4:
[rank0]: saved metadata: {'shape': torch.Size([8192]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 6:
[rank0]: saved metadata: {'shape': torch.Size([8192, 8192]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}

[hahmad2008]

@winglian @NanoCode012 Could you please check this?

[NanoCode012]

In current main, you can use that yaml. It should be aliased either way in the new commit, but it would be best to change to the new name.

[hahmad2008]

@NanoCode012 if i don't use the recent merged commit, and I set it to false like this:

model_kwargs:
   use_cache: False

WITH axolotl commit: git checkout -f 268543a3be20aa8de662761bb36645cbb5ff8f5b

and then I log these two dicts (self.model_config: , self.model_kwargs} at this line: i can see there is duplicate with the the same arg:

self.model_config: LlamaConfig {
     .......
  "use_cache": true,
}


self.model_kwargs: {'use_cache': False, 
.....
}

And this cause an issue:

self.model_config: LlamaConfig {
  "_name_or_path": "openlm-research/open_llama_3b_v2",
  "architectures": [
    "LlamaForCausalLM"
  ],
  "attention_bias": false,
  "attention_dropout": 0.0,
  "bos_token_id": 1,
  "eos_token_id": 2,
  "head_dim": 100,
  "hidden_act": "silu",
  "hidden_size": 3200,
  "initializer_range": 0.02,
  "intermediate_size": 8640,
  "max_position_embeddings": 2048,
  "mlp_bias": false,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 26,
  "num_key_value_heads": 32,
  "pad_token_id": 0,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-06,
  "rope_scaling": null,
  "rope_theta": 10000.0,
  "tie_word_embeddings": false,
  "torch_dtype": "float16",
  "transformers_version": "4.48.1",
  "use_cache": true,
  "vocab_size": 32000
}

self.model_kwargs: {'use_cache': False, 'device_map': None, 'torch_dtype': torch.float16, 'quantization_config': BitsAndBytesConfig {
  "_load_in_4bit": true,
  "_load_in_8bit": false,
  "bnb_4bit_compute_dtype": "float16",
  "bnb_4bit_quant_storage": "bfloat16",
  "bnb_4bit_quant_type": "nf4",
  "bnb_4bit_use_double_quant": true,
  "llm_int8_enable_fp32_cpu_offload": false,
  "llm_int8_has_fp16_weight": false,
  "llm_int8_skip_modules": null,
  "llm_int8_threshold": 6.0,
  "load_in_4bit": true,
  "load_in_8bit": false,
  "quant_method": "bitsandbytes"
}
, 'attn_implementation': 'flash_attention_2'}
`low_cpu_mem_usage` was None, now default to True since model is quantized.
[2025-02-08 20:37:19,460] [ERROR] [axolotl.load_model:1051] [PID:815] [RANK:0] LlamaForCausalLM.__init__() got an unexpected keyword argument 'use_cache'
Traceback (most recent call last):
  File "axolotl/src/axolotl/utils/models.py", line 1049, in load_model
    skip_move_to_device = self.build_model(qlora_fsdp)
  File "axolotl/src/axolotl/utils/models.py", line 821, in build_model
    self.model = self.AutoModelLoader.from_pretrained(
  File "/usr/local/lib/python3.10/dist-packages/transformers/models/auto/auto_factory.py", line 564, in from_pretrained
    return model_class.from_pretrained(
  File "/usr/local/lib/python3.10/dist-packages/transformers/modeling_utils.py", line 4090, in from_pretrained
    model = cls(config, *model_args, **model_kwargs)
TypeError: LlamaForCausalLM.__init__() got an unexpected keyword argument 'use_cache'

[NanoCode012]

Hey, one is for the model's config (LlamaConfig), the latter is for LlamaForCausalLM.from_pretrained.

Perhaps, we can take a step back? What do you want to achieve?

Your earlier question was how to pass to .from_pretrained which I informed was via model_kwargs. It seems that's not the solution as LlamaForCausalLM doesn't accept use_cache, and we default set use_cache=False in the LlamaConfig, so I'm not sure this is what you want to do.

Have you tried FSDP again?

[hahmad2008]

@NanoCode012 I tried to enable this use_reentrant: true which fix the error I got.

[NanoCode012]

Thank you for letting us know. I'm glad that solved it. Sorry for the roundabout convo we had, I lost track of the problem mid-way.

Could you mark that as the solution for future readers / update the main post with the solution ?

[hahmad2008]

Sure, just let me know how t do it? there is no option to mark a comment as solution

[NanoCode012]

I just marked it. It seems to only work in Q&A category. Thanks for following up!