[metax] modify some config (FlagOpen#740)

* modify config

* add config

* modify

---------

Co-authored-by: yaguang.wuyaguang <yaguang.wu@metax-tech.com>

training/metax/docker_image/megatron_core060/Dockerfile

@@ -30,6 +30,7 @@ ENV MHA_BWD_NO_ATOMIC_F64=1
 ENV MCCL_IB_GID_INDEX=1
 ENV NVTE_FLASH_ATTN=1
 ENV NVTE_FUSED_ATTN=0
+ENV MCCL_MAX_NCHANNELS=16
 
 ENV MCBLAS_CUSTOMIZED_CONFIG_PATH=/workspace/Megatron-LM_metax/mcblas_customized_config.yaml
 

training/metax/llama3_8B-megatron/config/config_C500x1x8.py

@@ -4,4 +4,4 @@
 theoryflops = -1
 megatron_path = "/workspace/Megatron-LM_metax" # need to be aligned with DockerFile. In NGCtorch, it's /workspace/ + Megatron-LM
 tensor_parallel = 1
-pipeline_parallel = 4
+pipeline_parallel = 2

training/metax/llama3_8B-megatron/config/training_adapter.sh

@@ -1,21 +1,22 @@
 # for 1*8 3150
-VENDOR_ARGS=" \
-    --transformer-impl local  \
-    --use-distributed-optimizer \
-    --use-mcore-models \
-    --use-flash-attn \
-    --pipline-num-layers-list 7 9 9 7
-"
-# for 4*8
 # VENDOR_ARGS=" \
 #     --transformer-impl local  \
 #     --use-distributed-optimizer \
 #     --use-mcore-models \
 #     --use-flash-attn \
-#     --attention-dropout 0.0 \
-#     --hidden-dropout 0.0 \
-#     --recompute-granularity full \
-#     --recompute-method block \
-#     --recompute-num-layers 1 \
-#     --recompute-num-layers-list 2 0
-# "
+#     --pipline-num-layers-list 7 9 9 7
+# "
+# for 4*8
+VENDOR_ARGS=" \
+    --transformer-impl local  \
+    --use-distributed-optimizer \
+    --use-mcore-models \
+    --use-flash-attn \
+    --attention-dropout 0.0 \
+    --hidden-dropout 0.0 \
+    --pipline-num-layers-list 16 16 \
+    --recompute-granularity full \
+    --recompute-method block \
+    --recompute-num-layers 1 \
+    --recompute-num-layers-list 5 0
+"