def fast_approx_rand(out: FloatTensorType) -> None: out = out.flatten() numel = out.numel() if numel < 1_000_003: torch.randn(numel, out=out) return t = torch.randn(1_000_003) excess = numel % 1_000_003 # Using just `-excess` would give bad results when excess == 0. out[:numel - excess].view(-1, 1_000_003)[...] = t out[numel - excess:] = t[:excess]
def store( self, entity: EntityName, part: Partition, embs: FloatTensorType, optim_state: Optional[bytes], ) -> None: key = "%s_%s" % (entity, part) futs = [] # embs sharded across all servers. num_clients = len(self._clients) flattened_embs = embs.flatten() numel_per_shard = self.shard_size // flattened_embs.element_size() # Space out the start of each partition's chunks of servers evenly. offset = (part * num_clients) // self.entities[entity].num_partitions for (idx, flattened_embs_shard) in enumerate( torch.split(flattened_embs, numel_per_shard)): client_idx = (idx + offset) % num_clients shard_key = key + f"__embs_{idx}" fut = self._async_store(client_idx, shard_key, flattened_embs_shard) futs.append(fut) # optim_state unsharded due to its smaller size. if optim_state is not None: client_idx = offset % num_clients optim_key = key + "__optim" optim_state_tensor = bytes_to_bytetensor(optim_state) optim_state_fut = self._async_store(client_idx, optim_key, optim_state_tensor) t0 = time.monotonic() for fut in futs: fut.result() t1 = time.monotonic() if optim_state is not None: optim_state_fut.result() t2 = time.monotonic() if self.log_stats: embs_size = embs.numel() * embs.element_size() embs_time = t1 - t0 embs_speed = embs_size / embs_time message = ( f"Stored (entity {entity}, partition {part}). Embs of size {embs_size} " f"bytes stored in {embs_time} seconds ({embs_speed:,.0f} B/s)." ) if optim_state is not None: optim_state_size = (optim_state_tensor.numel() * optim_state_tensor.element_size()) optim_state_time = t2 - t1 optim_state_speed = optim_state_size / optim_state_time message += ( f"optim_state of size {optim_state_size} bytes stored in " f"{optim_state_time} seconds ({optim_state_speed:,.0f} B/s)." ) logger.info(message)