Feature request: Support sharding parameters where first dimension is not divisible by 8

[mpatel31415]

🐛 Bug

With newest version of Docker image (tested on 2024-04-22) during training of Nous-Hermes-13b with thunder.jit with default executors training with FSDP + zero2 I get an error:

File "/opt/pytorch/lightning-thunder/thunder/core/baseutils.py", line 103, in check
[rank2]: raise exception_type(s())
[rank2]: RuntimeError: Current sharding requires the first dimension of the parameter 'lm_head.weight' (32001) to be divisible by the world size (8)

To Reproduce

1. Start container before every test:

docker run --gpus all --ipc=host --ulimit memlock=-1 --ulimit stack=67108864  --pull always --INTERNAL_ADDRESS:5005/dl/pytorch/update-scripts:pjnl-latest

Thunder
2. Inside container run:

torchrun --nproc-per-node=8 /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py --model_name Nous-Hermes-13b --compile thunder --distributed_mode fsdp --shard_mode zero2

Inductor:
2. Inside container run:

torchrun --nproc-per-node=8 /opt/pytorch/lightning-thunder/thunder/benchmarks/benchmark_litgpt.py --model_name Nous-Hermes-13b --compile inductor --distributed_mode fsdp --shard_mode zero2

This command works.

Expected behavior

We should be able to train Nous-Hermes-13b with FSDP + zero2 + Thunder.

Environment

As in the Docker image. The reproduction examples come from 8xA100 GPUs.

Output of nvidia-smi:

cc @carmocca @awaelchli @crcrpar

[mruberry]

triage review:

• updated title

Thanks @mpatel31415!