matmul dtensor experiments

[ClaudiaComito]

Due Diligence

• General:
  • title of the PR is suitable to appear in the Release Notes
• Implementation:
  • unit tests: all split configurations tested
  • unit tests: multiple dtypes tested
  • NEW unit tests: MPS tested (1 MPI process, 1 GPU)
  • benchmarks: created for new functionality
  • benchmarks: performance improved or maintained
  • documentation updated where needed

Description

Eventually we would like to implement NCCL or torch.distributed as alternative communication backends for optimized communication on multi-GPU.

A relatively straightforward way for selected operations would be to pipe relevant DNDarray operands into torch DTensors. Here's an example that works for matrix multiplication.

Note that torch.distr doesn't implement the -v and -w collective calls (Allgatherv, Alltoallv....) and therefore only supports same-size slices/chunks on each process. The DTensor backend is only available for equally distributed arrays on GPU.

More on tests and benchmarks in the comments.

Issue/s resolved: #

Changes proposed:

• add mesh initialization in communication.py
• add helper functions in linalg/basics.py (for now) to go from DNDarray to DTensor and vice versa - these should probably move to factories or the dndarray module or even a dedicated module
• ht.matmul is taking off to DTensor.matmul if the DNDarrays are on GPU and distributed equally.

Type of change

Memory requirements

Performance

Does this change modify the behaviour of other functions? If so, which?

yes / no