github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-04 02:26:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

[hotfix]: add pp sanity check and fix mbs arg (#5268)

Browse Source 
* fix: fix misleading mbs arg

* feat: add pp sanity check

* fix: fix 1f1b sanity check
This commit is contained in:
Wenhao Chen
2024-01-15 15:57:40 +08:00
committed by GitHub
parent c174c4fc5f
commit ef4f0ee854
5 changed files with 21 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
examples/language/llama2/benchmark.py
View File

@@ -74,8 +74,8 @@ def main():
parser.add_argument("--tp", type=int, default=1, help="Tensor parallel size")
parser.add_argument("--extra_dp", type=int, default=1, help="Extra data parallel size, used for Gemini")
parser.add_argument("--pp", type=int, default=1, help="Pipeline parallel size")
parser.add_argument("--mbs", type=int, default=1)
parser.add_argument("--zero", type=int, default=0)
parser.add_argument("--mbs", type=int, default=1, help="Micro batch size of pipeline parallel")
parser.add_argument("--zero", type=int, default=0, help="Zero Stage when hybrid plugin is enabled")
args = parser.parse_args()
colossalai.launch_from_torch({})
@@ -98,7 +98,13 @@ def main():
extra_dp_size=args.extra_dp,
)
elif args.plugin == "gemini_auto":
plugin = GeminiPlugin(placement_policy="auto", precision="bf16", warmup_non_model_data_ratio=args.warmup_ratio, tp_size=args.tp, extra_dp_size=args.extra_dp)
plugin = GeminiPlugin(
placement_policy="auto",
precision="bf16",
warmup_non_model_data_ratio=args.warmup_ratio,
tp_size=args.tp,
extra_dp_size=args.extra_dp,
)
elif args.plugin == "fsdp":
if use_empty_init:
plugin = TorchFSDPPlugin(
@@ -137,7 +143,7 @@ def main():
zero_stage=args.zero,
num_model_chunks=2,
enable_fused_normalization=torch.cuda.is_available(),
num_microbatches=args.mbs,
microbatch_size=args.mbs,
precision="bf16",
)
elif args.plugin == "3d_cpu":
@@ -147,7 +153,7 @@ def main():
zero_stage=args.zero,
cpu_offload=True,
enable_fused_normalization=torch.cuda.is_available(),
num_microbatches=args.mbs,
microbatch_size=args.mbs,
initial_scale=2**8,
precision="bf16",
)

2
examples/language/llama2/scripts/benchmark_70B/3d.sh
View File

@@ -14,4 +14,4 @@ cd ../..
export OMP_NUM_THREADS=8
colossalai run --nproc_per_node 8 --hostfile $HOSTFILE benchmark.py -c 70b -p 3d -g -x -b 8 --tp 4 --pp 2 --mbs 4
colossalai run --nproc_per_node 8 --hostfile $HOSTFILE benchmark.py -c 70b -p 3d -g -x -b 8 --tp 4 --pp 2 --mbs 1