use Environment Variable ONEFLOW_COMM_NET_IB_ENABLE

Classification/cnns/config.py

@@ -113,9 +113,6 @@ def get_parser(parser=None):
     parser.add_argument("--batch_size_per_device", type=int, default=64)
     parser.add_argument("--val_batch_size_per_device", type=int, default=8)
 
-    parser.add_argument(
-        "--use_rdma", type=str2bool, nargs="?", const=True, help="Use rdma.",
-    )
     parser.add_argument(
         "--nccl_fusion_threshold_mb",
         type=int,

Classification/cnns/of_cnn_train_val.py

@@ -64,9 +64,6 @@ if args.nccl_fusion_threshold_mb:
 if args.nccl_fusion_max_ops:
     flow.config.collective_boxing.nccl_fusion_max_ops(args.nccl_fusion_max_ops)
 
-if args.num_nodes > 1 and args.use_rdma:
-    flow.config.use_rdma(True)
-
 
 def label_smoothing(labels, classes, eta, dtype):
     assert classes > 0

Classification/cnns/train_fp16.sh

@@ -24,6 +24,7 @@ export PYTHONUNBUFFERED=1
 echo PYTHONUNBUFFERED=$PYTHONUNBUFFERED
 export NCCL_LAUNCH_MODE=PARALLEL
 echo NCCL_LAUNCH_MODE=$NCCL_LAUNCH_MODE
+export ONEFLOW_COMM_NET_IB_ENABLE=1
 
 python3 of_cnn_train_val.py \
      --train_data_dir=$DATA_ROOT/train \

LanguageModeling/GPT/README.md

@@ -173,7 +173,7 @@ python3 tools/launch_container.py \
 bash examples/pretrain_1n8d_2x4x1_16_1536x16.sh
 ```
 
-如果需要进一步扩充设备数量，由于单台机器适配的 GPU 设备数量有限，我们需要更多的物理机器 node。此时，我们需要配置 `--num-nodes` 和 `--node-ips` 参数，并且分别在每台机器上启动训练命令（`traning.py`）。同时在有 rdma 的环境中，可以开启 `--use-rdma` 来带来更佳的训练效率。以下是4机8卡下，各种并行方式混合的示例：
+如果需要进一步扩充设备数量，由于单台机器适配的 GPU 设备数量有限，我们需要更多的物理机器 node。此时，我们需要配置 `--num-nodes` 和 `--node-ips` 参数，并且分别在每台机器上启动训练命令（`traning.py`）。同时在有 rdma 的环境中，可以开启 `export ONEFLOW_COMM_NET_IB_ENABLE=1` 来带来更佳的训练效率。以下是4机8卡下，各种并行方式混合的示例：
 
 ```
 bash examples/distribute_pretrain_4n8d_2x4x4_512_2304x24.sh

LanguageModeling/GPT/examples/pretrain.sh

@@ -84,7 +84,7 @@ cmd+=" --multihead-attention-fusion"
 cmd+=" --fp16"
 
 if [[ ${num_nodes} -gt 1 ]]; then
-    cmd+=" --use-rdma"
+    export ONEFLOW_COMM_NET_IB_ENABLE=1
 fi
 
 if [[ ! -z "${ONEFLOW_GTP_PROFILE_FILE}" ]]; then

LanguageModeling/GPT/oneflow_gpt/config.py

@@ -635,11 +635,7 @@ def _add_misc_args(parser):
         action="store_true",
         help="open transformer layer profiler",
     )
-    group.add_argument(
-        "--use-rdma",
-        action="store_true",
-        help="Use rdma.",
-    )
+
     return parser
 
 

LanguageModeling/GPT/oneflow_gpt/training.py

@@ -47,8 +47,6 @@ def _init_config(args):
                 " please try other version."
             )
 
-    if args.use_rdma:
-        flow.config.use_rdma(True)
 
     flow.config.enable_legacy_model_io()
     flow.config.enable_model_io_v2(True)