How can I use only OSS in Pytorch lightning?

[CD21a]

Hello

I wanna use only OSS (zero1) in pytorch lightning, but pytorch lightning only support OSS + SDP (zero2).
How can i use only OSS?

[awaelchli]

Hi
You can do optimizer state sharding with deepspeed stage 1 (zero1):

trainer = Trainer(accelerator="gpu", devices=4, strategy="deepspeed_stage_1", precision=16)

[awaelchli]

If you don't want to use deepspeed, you can also just use the OSS optimizer in PyTorch: https://pytorch.org/tutorials/recipes/zero_redundancy_optimizer.html