Repository: microsoft/DeepSpeed
File: csrc/compile/z3.cpp (C++)

Code before:
                         c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem)
    {
        const DSParam& param = param_registry_->getParam(ds_id);
        const at::Tensor& ds_tensor = param.getDSTensor();

        if (symm_mem == nullptr) {
            // Support uneven shard sizes across ranks by padding to max shard size
            int world_size = process_group_->getSize();
            int rank = process_group_->getRank();

            int64_t local_count = ds_tensor.numel();

            // Gather local shard sizes from all ranks
            auto count_options = at::TensorOptions().dtype(at::kLong).device(at::kCUDA);
            at::Tensor local_count_tensor = torch::tensor({local_count}, count_options);
            std::vector<at::Tensor> all_counts(world_size);
            for (int i = 0; i < world_size; ++i) {
                all_counts[i] = torch::empty_like(local_count_tensor);
            }
            process_group_->allgather(all_counts, local_count_tensor)->wait();

            int64_t max_count = 0;
            std::vector<int64_t> host_counts(world_size);
            for (int i = 0; i < world_size; ++i) {
                host_counts[i] = all_counts[i].to(torch::kCPU).item<int64_t>();
                if (host_counts[i] > max_count) { max_count = host_counts[i]; }
            }

            // Prepare padded send buffer and gather buffer on AG stream
            at::Tensor send_buf;
            at::Tensor gather_tmp;
            {
                at::cuda::CUDAStreamGuard guard(ag_stream_);
                send_buf = torch::empty({max_count}, ds_tensor.options());
                // Copy real shard
                send_buf.index_put_({torch::indexing::Slice(0, local_count)}, ds_tensor.flatten(), true);
                // Zero-pad the tail if needed
                if (local_count < max_count) {
                    auto pad_len = max_count - local_count;
                    send_buf.index_put_({torch::indexing::Slice(local_count, max_count)},
                                        torch::zeros({pad_len}, ds_tensor.options()),
                                        true);
                }
                gather_tmp = torch::empty({static_cast<long>(world_size) * max_count}, ds_tensor.options());
            }

            ncclResult_t result = ncclAllGather(send_buf.data_ptr(),
                                                gather_tmp.data_ptr(),
                                                max_count,
                                                get_nccl_data_type(ds_tensor.scalar_type()),
                                                nccl_comm_,

Review comment (performance):
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.

Code after:
                         c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem)
    {
        const DSParam& param = param_registry_->getParam(ds_id);
        const at::Tensor& ds_tensor = param.getDSTensor();

        if (symm_mem == nullptr) {
            // Fast path: assume uniform shard sizes (ZeRO-3 partitions are padded to uniform size)
            const int world_size = process_group_->getSize();
            const int64_t shard_elems = ds_tensor.numel();

            // Perform all-gather directly into the pre-allocated padded output buffer
            // NCCL requires contiguous storage; use .contiguous() explicitly
            ncclResult_t result = ncclAllGather(ds_tensor.contiguous().data_ptr(),
                                                output_buf.data_ptr(),
                                                shard_elems,
                                                get_nccl_data_type(ds_tensor.scalar_type()),
                                                nccl_comm_,
                                                ag_stream_);

            if (result != ncclSuccess) { throw std::runtime_error("NCCL AllGather failed"); }
        } else {
            at::cuda::CUDAStreamGuard guard(ag_stream_);
            int world_size = process_group_->getSize();
            int rank = process_group_->getRank();

            at::Tensor local_buf =
                symm_mem->get_buffer(rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0);
            local_buf.copy_(ds_tensor, true);

            symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER);
            auto chunks = output_buf.flatten().chunk(world_size);
            for (int step = 0; step < world_size; step++) {
                int remote_rank = (rank - step + world_size) % world_size;
                auto src_buf = symm_mem->get_buffer(
                    remote_rank, ds_tensor.sizes(), ds_tensor.scalar_type(), 0);
                chunks[remote_rank].copy_(src_buf.flatten(), true);
            }
            symm_mem->barrier(0, TIMEOUT_SYMMETRIC_MEMORY_BARRIER);
        }

        param_registry_->registerGatheredParam(ds_id, output_buf);
        param_registry_->setValid(ds_id, true);
    }

    at::Tensor allgatherParam(long ds_id,
                              c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem)
    {
        const DSParam& param = param_registry_->getParam(ds_id);
        const at::Tensor& ds_tensor = param.getDSTensor();
        const int world_size = process_group_->getSize();
        const int64_t true_numel = static_cast<int64_t>(productDim(param.getShape()));