github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 17:17:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

[gemini] hotfix NaN loss while using Gemini + tensor_parallel (#5150)

Browse Source 
* fix

aaa

fix

fix

fix

* fix

* fix

* test ci

* fix ci

fix
This commit is contained in:
flybird11111
2023-12-08 11:10:51 +08:00
committed by GitHub
parent b397104438
commit 21aa5de00b
3 changed files with 59 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -72,6 +72,7 @@ def main():
parser.add_argument("--offload_optim_frac", type=float, default=0.0, help="Offload optim fraction. Only for gemini")
parser.add_argument("--offload_param_frac", type=float, default=0.0, help="Offload param fraction. Only for gemini")
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)
@@ -93,9 +94,11 @@ def main():
shard_param_frac=args.shard_param_frac,
offload_optim_frac=args.offload_optim_frac,
offload_param_frac=args.offload_param_frac,
tp_size=args.tp,
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)
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(