[zero] refactory ShardedOptimV2 init method (#416)

2025-10-25 18:13:56 +00:00 · 2022-03-15 10:45:55 +08:00
parent e79ea44247
commit 23ba3fc450
3 changed files with 32 additions and 16 deletions
--- a/tests/test_zero_data_parallel/test_sharded_optim_v2.py
+++ b/tests/test_zero_data_parallel/test_sharded_optim_v2.py
@@ -52,7 +52,6 @@ def run_dist(rank, world_size, port, cpu_offload, shard_strategy):
        optim = optimizer_class(model.parameters(), lr=lr)
        sharded_optim = ShardedOptimizerV2(zero_model,
                                           optimizer_class,
-                                           shard_strategy,
                                           cpu_offload=cpu_offload,
                                           initial_scale=2**5,
                                           lr=lr)