Multi-node checkpoint benchmark improvements

[MattIrv]

Several minor fixes/improvements for the multi-node checkpoint benchmark:

• Update the README with some missing instructions

• Run as a module rather than directly with Python so that imports work correctly

• Update WORLD_SIZE env var to NUM_NODES because PyTorch / PyTorch Lightning appear to use WORLD_SIZE for other purposes

• Set use_orig_params=False when creating the strategy, since this avoids the time-consuming step of flattening the optimizer state dict

  • This happened on each node, so I'm hopeful that this gives us a better trend for throughput when scaling the number of nodes
  • This also required updating the model's configure_optimizers to use self.trainer.model.parameters instead of self.model.parameters as in the default (see FSDP : The optimizer does not seem to reference any FSDP parameters. Lightning-AI/pytorch-lightning#17515)

• Make NUM_DEVICES a configurable environment variable

• Add the ability to configure the strategy with AdamW or SGD optimizer. AdamW appears to result in larger checkpoint files that are about 20% of the GPU memory utilization

• Tests pass

• Appropriate changes to documentation are included in the PR