rebase

Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Author: ananthsub

examples/llm/llama3_8b/conf/llama3_8b_pretrain_override_example.yaml

@@ -15,44 +15,44 @@
 # Example override file
 
 # To override a parameter, ensure the structure matches the ConfigContainer
-# and its sub-configurations (e.g., model_config, train_config, etc.)
+# and its sub-configurations (e.g., model, train, etc.)
 # Top-level ConfigContainer fields are dataclasses themselves
 
-model_config:
+model:
   tensor_model_parallel_size: 2
   pipeline_model_parallel_size: 2
   context_parallel_size: 1
   pipeline_dtype: torch.bfloat16
 
-train_config:
+train:
   train_iters: 500000
   global_batch_size: 256
   micro_batch_size: 4
   eval_interval: 1000
 
-optimizer_config:
+optimizer:
   lr: 0.00025
   min_lr: 0.000025
 
-scheduler_config:
+scheduler:
   lr_warmup_iters: 1000
 
-checkpoint_config:
+checkpoint:
   save_interval: 1000
 
-logger_config:
+logger:
   log_interval: 5
 
-dataset_config:
+dataset:
   num_workers: 4
 
-rng_config:
+rng:
   seed: 42
 
-ddp_config:
+ddp:
   grad_reduce_in_fp32: true
 
-profiling_config:
+profiling:
   use_nsys_profiler: false
   profile_step_start: 1000
   profile_step_end: 10000

examples/llm/llama3_8b/pretrain_llama3_8b.py

@@ -27,30 +27,12 @@
         $ python pretrain_llama3_8b.py --config-file my_custom_config.yaml
 
     Using CLI overrides only:
-        $ python pretrain_llama3_8b.py model_config.tensor_model_parallel_size=4 train_config.train_iters=100000
+        $ python pretrain_llama3_8b.py model.tensor_model_parallel_size=4 train.train_iters=100000
 
     Combining YAML and CLI overrides (CLI takes precedence):
         $ python pretrain_llama3_8b.py --config-file conf/my_config.yaml \
-        model_config.pipeline_dtype=torch.float16 \
-        train_config.global_batch_size=512
-
-    Adding new parameters with Hydra syntax:
-        $ python pretrain_llama3_8b.py +experiment_name=my_experiment +tags=[experiment,test]
-
-    Nested parameter overrides:
-        $ python pretrain_llama3_8b.py \
-        optimizer_config.lr=0.0001 \
-        scheduler_config.lr_warmup_iters=2000 \
-        checkpoint_config.save_interval=500
-
-    Multiple CLI overrides:
-        $ python pretrain_llama3_8b.py \
-            model_config.tensor_model_parallel_size=8 \
-            model_config.pipeline_model_parallel_size=1 \
-            train_config.train_iters=1000000 \
-            train_config.global_batch_size=1024 \
-            optimizer_config.lr=0.0003 \
-            dataset_config.num_workers=8
+        model.pipeline_dtype=torch.float16 \
+        train.global_batch_size=512
 
 Configuration Precedence:
     1. Base configuration from pretrain_config() recipe
@@ -127,16 +109,16 @@ def main() -> None:
 
     Examples of CLI usage:
         # Use default config with custom learning rate
-        python pretrain_llama3_8b.py optimizer_config.lr=0.0002
+        python pretrain_llama3_8b.py optimizer.lr=0.0002
 
         # Custom config file with additional overrides
-        python pretrain_llama3_8b.py --config-file my_config.yaml train_config.train_iters=50000
+        python pretrain_llama3_8b.py --config-file my_config.yaml train.train_iters=50000
 
         # Multiple overrides for distributed training
         python pretrain_llama3_8b.py \
-            model_config.tensor_model_parallel_size=4 \
-            model_config.pipeline_model_parallel_size=2 \
-            train_config.global_batch_size=512
+            model.tensor_model_parallel_size=4 \
+            model.pipeline_model_parallel_size=2 \
+            train.global_batch_size=512
     """
     args, cli_overrides = parse_cli_args()