sharded training

app.py

@@ -45,18 +45,19 @@ def main():
     model.load_state_dict(torch.load('model.pt'))
     optim = AdamW(model.parameters(), 2e-5)
 
-    bs = 24
-    kk = 128
+    bs = 1
+    kk = 2048
     dsx = C4X(kk+1)
     dlx = DataLoader(
         dsx,
         batch_size=bs, 
-        num_workers=8,
+        num_workers=4,
     )
 
     prog = tqdm(dlx, disable=not accelerator.is_main_process)
     
-    model, optim, dlx = accelerator.prepare(model, optim, dlx)
+    model = accelerator.prepare(model)
+    optim, dlx = accelerator.prepare(optim, dlx)
     
     optim.zero_grad()
     for i, batch in enumerate(prog):

default_config.yaml

@@ -0,0 +1,28 @@
+command_file: null
+commands: null
+compute_environment: LOCAL_MACHINE
+deepspeed_config: {}
+distributed_type: FSDP
+downcast_bf16: 'no'
+dynamo_backend: 'NO'
+fsdp_config:
+  fsdp_auto_wrap_policy: SIZE_BASED_WRAP
+  fsdp_backward_prefetch_policy: BACKWARD_PRE
+  fsdp_min_num_params: 2000
+  fsdp_offload_params: false
+  fsdp_sharding_strategy: 1
+  fsdp_state_dict_type: FULL_STATE_DICT
+gpu_ids: null
+machine_rank: 0
+main_process_ip: null
+main_process_port: null
+main_training_function: main
+megatron_lm_config: {}
+mixed_precision: 'no'
+num_machines: 1
+num_processes: 2
+rdzv_backend: static
+same_network: true
+tpu_name: null
+tpu_zone: null
+use_cpu: false