Fix

EasyLM/models/llama/llama_model.py

@@ -343,7 +343,9 @@ class LLaMAConfig(PretrainedConfig):
     @classmethod
     def get_tokenizer(cls, config, padding_side='left', truncation_side='right'):
         config = cls.get_tokenizer_config(config)
-        assert config.vocab_file != '' and config.pretrained_model_name_or_path != '', 'vocab_file or pretrained_model_name_or_path must be specified'
+        if config.vocab_file == '':
+            assert config.pretrained_model_name_or_path != '', 'vocab_file or pretrained_model_name_or_path must be specified'
+        
         if config.pretrained_model_name_or_path != '':
             tokenizer = AutoTokenizer.from_pretrained(
                 config.pretrained_model_name_or_path, 

pretrain_llama_7b.sh

@@ -8,7 +8,7 @@ export LIBTPU_INIT_ARGS='--xla_jf_spmd_threshold_for_windowed_einsum_mib=0 --xla
 
 
 python3 -m EasyLM.models.llama.llama_train \
-    --initialize_jax_distributed=True \
+    --jax_distributed.initialize_jax_distributed=True \
     --mesh_dim='1,-1,4' \
     --dtype='bf16' \
     --total_steps=1000000 \
@@ -20,7 +20,7 @@ python3 -m EasyLM.models.llama.llama_train \
     --update_llama_config='' \
     --load_dataset_state='' \
     --load_checkpoint='' \
-    --tokenizer.vocab_file='tokenizer.model' \
+    --tokenizer.pretrained_model_name_or_path='./' \
     --optimizer.type='lion' \
     --optimizer.lion_optimizer.weight_decay=1.0 \
     --optimizer.lion_optimizer.lr=3e-5 \