Support LLaMA training through SPMD

[jonb377]

Code change to enable training LLaMA through SPMD.

[alanwaketan]

LGTM.

[alanwaketan]

The new change LGTM.

[alanwaketan]

@jonb377 As discussed in the chat, can we have a new branch called llama2-google-next-training and then land your change there?

[jonb377]

Sounds good, I've changed the base to llama2-google-next-training. I'll merge this now and follow up with tensor parallelism.

[alanwaketan]

@jonb377 Wondering if it makes sense for you to share the command you use to run LLaMA 2 on the repo? Maybe on the "How to Run HF Transformers with SPMD" doc?

[jonb377]

Sure! It's essentially the same training command, you just need to update your model config file to be llama, e.g.:

{
  "_name_or_path": "/tmp/home/config.json",
  "bos_token_id": 1,
  "eos_token_id": 2,
  "hidden_act": "silu",
  "hidden_size": 4096,
  "initializer_range": 0.02,
  "intermediate_size": 11008,
  "max_position_embeddings": 2048,
  "model_type": "llama",
  "num_attention_heads": 32,
  "num_hidden_layers": 32,
  "num_key_value_heads": 32,
  "pad_token_id": 0,
  "pretraining_tp": 1,
  "rms_norm_eps": 1e-06,
  "rope_scaling": null,
  "tie_word_embeddings": false,
  "transformers_version": "4.32.0.dev0",
  "use_cache": true,
  "vocab_size": 32000
}

[alanwaketan]

Thanks @jonb377! It will be great if you can share your model configs in a gs bucket just like what Alex did.

[jonb377]

http://bigstore/hf-train-config, I'll also link in the doc

[alanwaketan]

Thanks, Jon!