writeback=True and rank0_only=True currently raises an error. We should be able to implement this setting by having rank 0 broadcast the unsharded FlatParameter to all ranks before writing back.

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501