| 117 | } |
| 118 | |
| 119 | at::Tensor allgatherParam(long ds_id, |
| 120 | std::optional<at::ScalarType> dtype, |
| 121 | c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem) |
| 122 | { |
| 123 | const DSParam& param = param_registry_->getParam(ds_id); |
| 124 | const at::Tensor& ds_tensor = param.getDSTensor(); |
| 125 | const int world_size = process_group_->getSize(); |
| 126 | const int64_t true_numel = static_cast<int64_t>(productDim(param.getShape())); |
| 127 | const int64_t padded_per_rank = (true_numel + world_size - 1) / world_size; |
| 128 | const int64_t padded_numel = static_cast<int64_t>(world_size) * padded_per_rank; |
| 129 | at::ScalarType target_dtype = dtype ? dtype.value() : ds_tensor.scalar_type(); |
| 130 | |
| 131 | if (param_registry_->isValid(ds_id)) { |
| 132 | // Return a view sliced to the true size with the original shape |
| 133 | // |
| 134 | // Persistent params are gathered in their original dtype which may |
| 135 | // be different from the requested. |
| 136 | auto base = param_registry_->getGatheredParam(ds_id); |
| 137 | return base.flatten() |
| 138 | .to(target_dtype) |
| 139 | .index({torch::indexing::Slice(0, true_numel)}) |
| 140 | .view(param.getShape()); |
| 141 | } |
| 142 | |
| 143 | at::Tensor output_buf; |
| 144 | if (param_registry_->hasGatheredParam(ds_id)) { |
| 145 | auto existing = param_registry_->getGatheredParam(ds_id); |
| 146 | if (existing.defined() && existing.numel() == padded_numel) { output_buf = existing; } |
| 147 | } |
| 148 | if (!output_buf.defined()) { |
| 149 | at::cuda::CUDAStreamGuard guard(ag_stream_); |
| 150 | output_buf = torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); |
| 151 | } |
| 152 | |
| 153 | assert(hasKey(ag_comp_done_events_, ds_id)); |
| 154 | ag_comp_done_events_[ds_id]->record(); |
| 155 | ag_comp_done_events_[ds_id]->block(ag_stream_); |
| 156 | |
| 157 | launchAllGather(output_buf, ds_id, symm_mem); |
| 158 | |
| 159 | ag_comm_done_events_[ds_id]->record(ag_stream_); |
| 160 | // Return a view of the gathered padded buffer matching the true param shape |
| 161 | return output_buf.flatten() |
| 162 | .index({torch::indexing::Slice(0, true_numel)}) |
| 163 | .view(param.getShape()); |
| 164 | } |
| 165 | |
| 166 | void prefetchParamsFused(const std::vector<long>& ds_ids, |
| 167 | const std::optional<std::vector<at::ScalarType>> dtypes, |
no test coverage detected