refactored initial implementation to use torch fsdp api and adding new tests

docs/source/multi_gpu_training.mdx

@@ -152,12 +152,21 @@ optimizer before `accelerator.prepare()`.
 ### FSDP checkpoints
 
 LeRobot gathers the full state dict across all ranks and the main process writes it as a single
-`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two thigs to look out for:
+`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for:
 
-- With mixed precision, (`bf16`/`fp16`) FSDP keeps an fp32 master copy, so the checkpoint is fp32
-  (~2× the bf16 size on disk) and is cast back to the policy dtype on load.
-- **Optimizer state is not saved under FSDP**, so **resume-from-checkpoint is not supported**.
-  Saved weights are fully usable for evaluation and fine-tuning.
+- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master
+  copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently
+  with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific
+  caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting,
+  so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU
+  first, or cast `model.safetensors` to the deployment dtype offline.
+- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved
+  alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json`
+  format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`.
+  Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can
+  resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only
+  sample-exact when the world size and batch size match the original run (a warning is logged
+  otherwise); the optimizer/model state itself is unaffected.
 
 ## Notes