Browse Source

[NFC] polish applications/Chat/benchmarks/benchmark_opt_lora_dummy.py code style (#5225)

fix/format
Ofey Chan 6 months ago committed by GitHub
parent
commit
c7d68b2c2c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
  1. 6
      applications/Chat/benchmarks/benchmark_opt_lora_dummy.py

6
applications/Chat/benchmarks/benchmark_opt_lora_dummy.py

@ -76,9 +76,11 @@ def main(args):
if args.strategy == "ddp": if args.strategy == "ddp":
strategy = DDPStrategy() strategy = DDPStrategy()
elif args.strategy == "colossalai_gemini": elif args.strategy == "colossalai_gemini":
strategy = GeminiStrategy(placement_policy="static",initial_scale=2**5) strategy = GeminiStrategy(placement_policy="static", initial_scale=2**5)
elif args.strategy == "colossalai_gemini_cpu": elif args.strategy == "colossalai_gemini_cpu":
strategy = GeminiStrategy(placement_policy="static", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5) strategy = GeminiStrategy(
placement_policy="static", offload_optim_frac=1.0, offload_param_frac=1.0, initial_scale=2**5
)
elif args.strategy == "colossalai_zero2": elif args.strategy == "colossalai_zero2":
strategy = LowLevelZeroStrategy(stage=2, placement_policy="cuda") strategy = LowLevelZeroStrategy(stage=2, placement_policy="cuda")
elif args.strategy == "colossalai_zero2_cpu": elif args.strategy == "colossalai_zero2_cpu":

Loading…
Cancel
Save