Feature/transformer sequence sharding

[japols]

This PR adds a new sharding strategy shard_sequence for the transformer processor.

The current implementation (shard_heads) alternates between sharding across the sequence to sharding across heads for the sliding window attention mechanism. This requires two all-to-all communication steps per layer.

The shard_sequence strategy simplifies this process by keeping a sequence shard on each GPU and computing the sliding window attention locally. This requires a halo communication to exchange overlapping window segments (halos) between neighboring sequence shards.

Instead of 2 all-to-all communication steps per layer, the halo exchange only requires a single point-to-point communication between neighbouring GPUs, reducing communication time and improving scalability of model sharding across multiple GPUs.

The following benchmarking results show that using a 2 neighbor all-to-all (orange) is the best communication strategy to implement the halo exchange which consistently outperforms the old head-sharding strategy (blue):

This is an isolated fwd+bwds pass of 16 transformer layers with o96 input shapes, 1024 channels.

For a full training run on n320, o96 hidden we get the following increases in throughput (aligning with the benchmark results):

GPUs/Model	sharding strategy	avg time/batch (s)
2	shard_heads	1.38495
2	shard_sequence	1.29771
4	shard_heads	0.72034
4	shard_sequence	0.69254

mlflow

[FussyDuck]

All committers have signed the CLA.

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.85%. Comparing base (225315e) to head (063844b).

Additional details and impacted files

@@           Coverage Diff            @@
##           develop      #90   +/-   ##
========================================
  Coverage    99.85%   99.85%           
========================================
  Files           23       23           
  Lines         1374     1374           
========================================
  Hits          1372     1372           
  Misses           2        2           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.