out of memory for continuing pretraining llama3-8B

[ckzbullbullet]

I am trying to use the framework to continue pretraining llama3-8B. I have converted the HF checkpoint into nanotron format and the generated tokens seem reasonable.

I use the following setting to train the model but I got the OOO with 8 GPUs. I tried HF accelerate with deepspeed zero-1 and flash-attention and it worked well previously:

model:
ddp_bucket_cap_mb: 25
dtype: bfloat16
init_method:
std: 0.025
make_vocab_size_divisible_by: 1
model_config:
bos_token_id: 128000
eos_token_id: 128001
hidden_act: silu
hidden_size: 4096
initializer_range: 0.02
intermediate_size: 14336
is_llama_config: true
max_position_embeddings: 8192
num_attention_heads: 32
num_hidden_layers: 32
num_key_value_heads: 8
pad_token_id: null
pretraining_tp: 1
rms_norm_eps: 1.0e-05
rope_scaling: null
tie_word_embeddings: false
use_cache: true
vocab_size: 128256
optimizer:
accumulate_grad_in_fp32: true
clip_grad: 1.0
learning_rate_scheduler:
learning_rate: 8.0e-07
lr_decay_starting_step: null
lr_decay_steps: 100
lr_decay_style: cosine
lr_warmup_steps: 100
lr_warmup_style: linear
min_decay_lr: 1.0e-05
optimizer_factory:
adam_beta1: 0.9
adam_beta2: 0.999
adam_eps: 1.0e-08
name: adamW
torch_adam_is_fused: true
weight_decay: 0.1
zero_stage: 1
parallelism:
dp: 8
expert_parallel_size: 1
pp: 1
pp_engine: 1f1b
tp: 1
tp_linear_async_communication: true
tp_mode: REDUCE_SCATTER
profiler: null
tokenizer:
tokenizer_max_length: null
tokenizer_name_or_path: meta-llama/Meta-Llama-3-8B
tokenizer_revision: null
tokens:
batch_accumulation_per_replica: 1
limit_test_batches: 0
limit_val_batches: 0
micro_batch_size: 1
sequence_length: 8192
train_steps: 200
val_check_interval: -1

[zzhhjjj]

Hi,

Thanks for your question. Here is my solution:

general:
  benchmark_csv_path: null
  consumed_train_samples: null
  ignore_sanity_checks: true
  project: test
  run: llama
  seed: 42
  step: 1
lighteval: null
logging:
  iteration_step_info_interval: 1
  log_level: info
  log_level_replica: info
model:
  ddp_bucket_cap_mb: 25
  dtype: bfloat16
  init_method:
    std: 0.2
  make_vocab_size_divisible_by: 1
  model_config:
    bos_token_id: 128000
    eos_token_id: 128001
    hidden_act: silu
    hidden_size: 4096
    initializer_range: 0.02
    intermediate_size: 14336
    is_llama_config: true
    max_position_embeddings: 8192
    num_attention_heads: 32
    num_hidden_layers: 32
    num_key_value_heads: 8
    pad_token_id: null
    pretraining_tp: 1
    rms_norm_eps: 1.0e-05
    rope_scaling: null
    tie_word_embeddings: false
    use_cache: true
    vocab_size: 128256
optimizer:
  accumulate_grad_in_fp32: true
  clip_grad: 1.0
  learning_rate_scheduler:
    learning_rate: 0.0003
    lr_decay_starting_step: null
    lr_decay_steps: 198
    lr_decay_style: cosine
    lr_warmup_steps: 2
    lr_warmup_style: linear
    min_decay_lr: 1.0e-05
  optimizer_factory:
    adam_beta1: 0.9
    adam_beta2: 0.95
    adam_eps: 1.0e-08
    name: adamW
    torch_adam_is_fused: true
  weight_decay: 0.01
  zero_stage: 0
parallelism:
  dp: 2
  expert_parallel_size: 1
  pp: 1
  pp_engine: 1f1b
  tp: 4
  tp_linear_async_communication: true
  tp_mode: REDUCE_SCATTER
profiler: null
tokenizer:
  tokenizer_max_length: null
  tokenizer_name_or_path: meta-llama/Meta-Llama-3-8B
  tokenizer_revision: null
tokens:
  batch_accumulation_per_replica: 32
  limit_test_batches: 0
  limit_val_batches: 0
  micro_batch_size: 1
  sequence_length: 8192
  train_steps: 30
  val_check_interval: -1

You can also try tp=8, dp=1.

Then should set theta = 500000.0 for RotaryEmbedding and interleaved = False.

Please let me know if you have any other questions

[ckzbullbullet]

Thanks for the reply.
I have implemented loading pretrained weights from llama-8b and it seems work.
Actually, my question is that when I use nanotron with pure DP (dp=8,pp=1,tp=1), I got OOO even if I use micro-batch-size=1. However, I did not get OOO when I use accelerate + deepspeed ZERO1 with the similar settings.
What would be the possible reason that nanotron uses more memory for pure DP? Also, if I want to implement that, where I should look into?

[zzhhjjj]

If my understanding is correct, ZERO1 is not pure DP, that's why you got the error.

Some links:
https://huggingface.co/transformers/v4.12.5/parallelism.html
https://blog.csdn.net/baoyan2015/article/details/136820078

[ckzbullbullet]

Sorry for the confusion.
I was comparing nanotron with the setting "dp=8,pp=1,tp=1,zero=1" with deepspeed zero1 without pp&tp.
I actually found that in deepspeed, I set the gradient_checkpointing=True to save memory during training but I did not find any notes about gradient_checkpoint in nanotron.

In the doc, I found the "activation checkpointing" with a decorator "@checkpoint_method(attr_name="do_checkpoint")". I applied it on LLaMAModel and I got the following error:
"File "/home/kezhen/Documents/nanotron/src/nanotron/parallel/pipeline_parallel/engine.py", line 56, in forward
[rank3]: assert output["loss"].requires_grad"

I have the following question:

1. Does nanotron support gradient_checkpointing right now?
2. Is the activation checkpointing similar with gradient checkpointing?
3. What is the correct way to use it?

[zzhhjjj]

Hello,

1,2 Yes, @checkpoint_method is the same as gradient checkpointing.
3 I didn't have enough time to check the code, but you can set use_reentrant=False to fix the error. Your observation is correct, even then there is still an OOO error, I will check the reason later when I have time.
https://pytorch.org/docs/stable/checkpoint.html

Thanks a lot for your questions.