update dist_ckpt format to zarr for PP>1 issue (#440)

Signed-off-by: paul-gibbons <[email protected]>
NVIDIA · Nov 18, 2024 · 64833d9 · 64833d9
1 parent 597f83a
commit 64833d9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/launcher_scripts/conf/fine_tuning/neva/mixtral_8x7b_instruct.yaml b/launcher_scripts/conf/fine_tuning/neva/mixtral_8x7b_instruct.yaml
@@ -91,6 +91,7 @@ model:
   # LLM configs
   # use GPTModel from megatron.core
   mcore_gpt: True
+  dist_ckpt_format: 'zarr'
   moe_grouped_gemm: false
   moe_token_dispatcher_type: alltoall
   moe_aux_loss_coeff: 0.01

diff --git a/launcher_scripts/conf/training/neva/mixtral_8x7b_instruct.yaml b/launcher_scripts/conf/training/neva/mixtral_8x7b_instruct.yaml
@@ -90,6 +90,7 @@ model:
   # LLM configs
   # use GPTModel from megatron.core
   mcore_gpt: True
+  dist_ckpt_format: 'zarr'
 
   moe_grouped_gemm: false
   moe_token_dispatcher_type: alltoall