Cannot save checkpoint while train.py on single GPU

[tiendung]

I got following error at utils.py while saving checkpoint while pre-training on single GPU. Any hint how should I fix it? Thanks.

state_dict = model._forward_module.state_dict()                            

NotImplementedError: offload_to_cpu=True and NO_SHARD is not supported yet

[lantiga]

@awaelchli is this something related to the current set of fixes or should we configure FSDP differently?

[lantiga]

@tiendung fix landed thanks to @awaelchli, can you give it a shot now?

[awaelchli]

FYI the fix just gets rid of the error, but to run training on a single GPU, you will need a lot of memory, and checkpionting can result in OOM. If you have trouble, you can create a smaller model by setting config.n_layers and config.n_embd for example.