[OSS] More flexible parameter handling in the low tensor case

[blefaudeux]

🚀 Feature

Have a better fallback than asserting when there are not enough tensors to shard

Motivation

#406 rightfully asserts for now when there are not enough tensors to shard, but in truth we could have a more graceful fallback in that case, not all ranks will contribute to the optim update but that's typically fine and not really a perf limitation

Pitch

When not enough tensors, just make some ranks data-parallel contributors and skip the sharding for them.

Alternatives

Current state which imposes a hard limitation on ranks/number of tensors in the model

Additional context

cc @kamo-naoyuki

[blefaudeux]

fixed