| 164 | } |
| 165 | |
| 166 | void prefetchParamsFused(const std::vector<long>& ds_ids, |
| 167 | const std::optional<std::vector<at::ScalarType>> dtypes, |
| 168 | c10::intrusive_ptr<c10d::symmetric_memory::SymmetricMemory> symm_mem) |
| 169 | { |
| 170 | std::vector<std::tuple<long, std::optional<at::ScalarType>>> invalid_params; |
| 171 | for (int i = 0; i < ds_ids.size(); i++) { |
| 172 | if (!param_registry_->isValid(ds_ids[i])) { |
| 173 | auto dtype = dtypes ? dtypes.value()[i] : std::optional<at::ScalarType>(); |
| 174 | invalid_params.push_back(std::make_tuple(ds_ids[i], dtype)); |
| 175 | } |
| 176 | } |
| 177 | |
| 178 | std::unordered_map<long, at::Tensor> output_bufs; |
| 179 | for (const auto& [ds_id, dtype] : invalid_params) { |
| 180 | const DSParam& param = param_registry_->getParam(ds_id); |
| 181 | const at::Tensor& ds_tensor = param.getDSTensor(); |
| 182 | const int world_size = process_group_->getSize(); |
| 183 | const int64_t shard_elems = ds_tensor.numel(); |
| 184 | const int64_t padded_numel = static_cast<int64_t>(world_size) * shard_elems; |
| 185 | |
| 186 | if (param_registry_->hasGatheredParam(ds_id)) { |
| 187 | auto existing = param_registry_->getGatheredParam(ds_id); |
| 188 | if (existing.defined() && existing.numel() == padded_numel) { |
| 189 | output_bufs[ds_id] = existing; |
| 190 | continue; |
| 191 | } |
| 192 | } |
| 193 | auto target_dtype = dtype ? dtype.value() : ds_tensor.scalar_type(); |
| 194 | output_bufs[ds_id] = |
| 195 | torch::empty({padded_numel}, ds_tensor.options().dtype(target_dtype)); |
| 196 | } |
| 197 | |
| 198 | for (const auto& [ds_id, _] : invalid_params) { |
| 199 | ag_comp_done_events_[ds_id]->record(); |
| 200 | ag_comp_done_events_[ds_id]->block(ag_stream_); |
| 201 | } |
| 202 | |
| 203 | ncclGroupStart(); |
| 204 | for (const auto& [ds_id, _] : invalid_params) { |
| 205 | assert(hasKey(output_bufs, ds_id)); |
| 206 | launchAllGather(output_bufs.at(ds_id), ds_id, symm_mem); |
| 207 | } |
| 208 | ncclGroupEnd(); |
| 209 | |
| 210 | for (const auto& [ds_id, _] : invalid_params) { |
| 211 | ag_comm_done_events_[ds_id]->record(ag_stream_); |
| 212 | } |
| 213 | } |
| 214 | |
| 215 | void releaseParam(long ds_id, long n_users) |
| 216 | { |
no test coverage detected