DeepCompile ZeRO-3: robust allgather for uneven shards; fix profiling…

[juyterman1000]

… meta key (max_mem)

[sfc-gh-truwase]

@juyterman1000 can you please address the formatting issue using https://github.com/deepspeedai/DeepSpeed/blob/master/CONTRIBUTING.md#prerequisites

[eternalNight]

Could you elaborate more on when ds_tensor.numel() of the same paramter can differ on different ranks? I think padding is already taken into account when the parameter is partitioned among the ranks (ref: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/runtime/zero/partition_parameters.py#L1664)

In case partition sizes do vary across ranks, can we fix that in partition_parameters.py to avoid synchronous communication here? launchAllGather() is on the critical path, so synchronous allgather can hurt performance.

[juyterman1000]

Thanks for the sharp catch., I’ve removed the synchronous size‑allgather from the hot path in launchAllGather() and now use a fixed‑count NCCL allgather, trimming any end padding to the true param size. To keep things better without paying the runtime cost, I added a one‑time registration‑time assertion that shard sizes match across ranks, if there’s ever a mismatch, we’ll catch it at source rather than synchronize in the critical path. Changes are in the updated PR.

[eternalNight]

I think we can further optimize the code by making allgatherParam() allocating a buffer with padding in the first place (today it allocates a buffer of ds_shape which is the true size of the gathered parameter). With that we don't need any additional memcpy or GPU memory allocation/deallocation. Instead we can slice the gathered output_buf before returning it. My understanding is that torch can correctly track the refcount to the underlying buffers even living tensors use only part of them, but correct me if I'm wrong.

[sfc-gh-truwase]

@eternalNight thanks for the suggestion. @juyterman1000 if you agree with this, do you want to address in a follow up PR? A benefit of a follow up PR is that it could document the perf benefit of the optimization separately from functionality.

[tohtana]

Hi @juyterman1000,

Thank you for the PR! As some changes are unclear to me, can you explain a bit more?
You now added an assertion to ensure the even sharding, which totally makes sense to me. Do we still need the changes launchAllGather? The additional memory allocation and copy might cause the significant overhead in some cases.

[juyterman1000]

@eternalNight Yes, we can allocate a buffer sized to world_size * shard_elems up front and slice it to the true size on return. PyTorch views hold a reference to the underlying storage; returning a sliced view does not break refcounting. We can cache the padded buffer per param to avoid repeat allocations. @sfc-gh-truwase . Agreed on the follow up.I’ll include micro-benchmarks showing the removal of 1 alloc + 1 memcpy per all-gather and any other gains. @tohtana With the even-sharding assertion in place, we don’t need the extra copy logic in launchAllGather(). We can issue a direct AllGather with a uniform shard ,elems count into the padded buffer and return a view of the first true_numel elements reshaped to the original param shape. The symmetric-memory path will stay as-is. This is to avoid additional copy overhead.

[tohtana]

@juyterman1000 The path using symmetric memory is experimental and not well optimized. So we need to keep non-symmetric memory path as the choice for the best performance.
If the allocation and copy are for uneven partitioning and the assertion block such an uneven partitioning, why can't we remove them?

[juyterman1000]

@juyterman1000 can you please address the formatting issue using https://github.com/deepspeedai/DeepSpeed/blob/master/CONTRIBUTING.md#prerequisites

Thanks for the check. I followed the formatting prerequisites in the contributing guide and ran the full pre-commit suite. I’ve pushed the updates.

[eternalNight]

Why replacing .contiguous() with .flatten()? .contiguous() makes sure that the underlying storage is contiguous which is required by nccl. .flatten() is a view-change and does not guarantee that.

Note: I believe the sharded tensors are already contiguous as they are already defragmented by DeepSpeedZeroOptimizer_Stage3.defragment(), but adding a .contiguous() does not hurt anyway and may help later when the layout of sharded tensors is changed.

[eternalNight]

I'm not sure when isValid(ds_id) is false while hasGatheredParam(ds_id) is true. They are both set at the end of launchAllGather(), and releasing a gathered param will unset the valid flag in unregisterGatheredParam().