Integrate changes for latest Distconv

[michaelmckinsey1]

LBANN/DistConv#13 Made some breaking changes

(1) ps.num_shards is now a tuple, instead of an int

[rank3]:   File "/usr/WS1/mckinsey/bp_scaffold-riken/ScaFFold/ScaFFold/utils/trainer.py", line 300, in train
[rank3]:     if images.size(0) < ps.num_shards:
[rank3]:        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank3]: TypeError: '<' not supported between instances of 'int' and 'tuple'

(2) mesh_dim_names format is different

[rank1]:   File "/usr/WS1/mckinsey/bp_scaffold-riken/ScaFFold/ScaFFold/worker.py", line 228, in main
[rank1]:     trainer.train()
[rank1]:   File "/usr/WS1/mckinsey/bp_scaffold-riken/ScaFFold/ScaFFold/utils/trainer.py", line 442, in train
[rank1]:     device_mesh=ps.device_mesh["dc"],
[rank1]:                 ~~~~~~~~~~~~~~^^^^^^
[rank1]:   File "/usr/WS1/mckinsey/bp_scaffold-riken/ScaFFold/.venvs/scaffoldvenv-tuo/lib/python3.11/site-packages/torch/distributed/device_mesh.py", line 730, in __getitem__
[rank1]:     slice_mesh_dims = _mesh_resources._get_slice_mesh_dims(
[rank1]:                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/usr/WS1/mckinsey/bp_scaffold-riken/ScaFFold/.venvs/scaffoldvenv-tuo/lib/python3.11/site-packages/torch/distributed/device_mesh.py", line 315, in _get_slice_mesh_dims
[rank1]:     raise KeyError(
[rank1]: KeyError: "Invalid mesh_dim_names ('dc',) specified. Valid mesh_dim_names are ['ddp', 'dc0']."

• Also on this - shouldn't the option be dc2? As our current default config sets shard_dim: 2, but maybe this is not being configured properly -> another note, maybe shard_dim: 2 should be shard_dim: 1 as @PatrickRMiles found this is the fastest configuration -> So all in all, this should probably be dc1

[jvwilliams23]

Hey @michaelmckinsey1, regarding point 2, we could change it such that the string is "dc2" instead of "dc0". I didn't realise people accessed this stuff from outside of distconv. @fiedorowicz1 what do you think?

[fiedorowicz1]

I agree it would be reasonable to rename those strings to match the actual shard dim rather so that they are more meaningful externally. Would that be best for you @michaelmckinsey1? If so, we can open a new PR and get it merged quickly.

Regarding setting shard_dim = 1 do you mean you are sharding along channel dimension?

[michaelmckinsey1]

To close the loop on this - We had a discussion and @fiedorowicz1 clarified some of our confusion with our usage of shard_dim, this parameter should be limited to the spatial dimensions (NC)DHW, independent of channels first/last. So we should only be setting shard_dim=[2-4] (DHW).

Additionally, we agree renaming the string to match the shard dim makes the most sense, thanks.