sharded

mosaicml · Jun 11, 2024 · 2796140 · 2796140
1 parent a255a6f
commit 2796140
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tests/trainer/test_fsdp_checkpoint.py b/tests/trainer/test_fsdp_checkpoint.py
@@ -935,6 +935,7 @@ def test_fsdp_partitioned_state_dict_load(
     if use_hsdp:
         fsdp_config = FSDPConfig(
             sharding_strategy='HYBRID_SHARD',
+            state_dict_type='sharded', 
             sharded_ckpt_prefix_dir='ba{batch}',
             data_parallel_shard_degree=data_parallel_shard_degree,
             data_parallel_replicate_degree=world_size // data_parallel_shard_degree,