def _broadcast_state_dict(self) -> None: """Broadcast this rank's state shard, discard others""" # Default to CPU space to gain some memory headroom local_cpu_state = recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu")) for rank in range(self.world_size): if rank == self.rank: # Send the state to the reference replica logging.debug( "Sending the sharded optimizer state to the reference replica from rank %s", rank, ) dist.broadcast_object_list([local_cpu_state], src=self.global_rank, group=self.group) else: global_rank = self.get_global_rank(self.group, rank) # Discard this tensor/rank, broadcast necessary for syncing and because NCCL does not support gather dist.broadcast_object_list([0], src=global_rank, group=self.group)
def _broadcast_processed_optim_state_dict( processed_optim_state_dict: Optional[Dict[str, Any]], rank: int, group, device: torch.device, ) -> Dict[str, Any]: """ Broadcasts the processed optimizer state dict from rank 0 to all ranks. Args: processed_optim_state_dict (Optional[Dict[str, Any]]): The flattened optimizer state dict with positive-dimension tensor states replaced with metadata if on rank 0; ignored otherwise. device (torch.device): Device to move zero-dimension tensors post- broadcast. Returns: Dict[str, Any]: The processed optimizer state dict. """ # Broadcast the two data structures rank 0 to all ranks obj_list = [processed_optim_state_dict] if rank == 0 \ else [None] dist.broadcast_object_list(obj_list, src=0, group=group) processed_optim_state_dict = obj_list[0] # type: ignore[assignment] assert processed_optim_state_dict is not None # Move zero-dimension tensors to `device` for param_state in processed_optim_state_dict["state"].values(): for state_name, value in param_state.items(): if _is_zero_dim_tensor(value): param_state[state_name] = value.to(device) return processed_optim_state_dict
def test_broadcast_object_list(self): val = 99 if self.rank == 0 else None object_list = [val] * dist.get_world_size() # TODO test with broadcast_object_list's device argument dist.broadcast_object_list(object_list=object_list) self.assertEqual(99, object_list[0])
def _collect_sharded_states(self) -> List[Dict[str, Any]]: """Collect all the state shards, in CPU memory.""" all_states = [] for rank in range(self.world_size): if rank == self.rank: logging.debug("Saving self state") all_states.append( recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))) # Sync with other replicas dist.broadcast_object_list([0], src=self.global_rank, group=self.group) else: # Fetch the optim state from the other replicas global_rank = self.get_global_rank(self.group, rank) replica_state = [0] dist.broadcast_object_list(replica_state, src=global_rank, group=self.group) all_states.append( recursive_copy_to_device(replica_state[0], non_blocking=True, device=torch.device("cpu"))) logging.debug("State from rank %s received", rank) return all_states
def sync_object_ranks(something_to_sync: Any, reference_rank: int, device: torch.device) -> Any: package = [something_to_sync] dist.broadcast_object_list(package, src=reference_rank, group=dist.group.WORLD) package_sync = package[0] return package_sync
def _broadcast_full_osd(self, full_osd: Dict[str, Any], group=None): """Broadcasts the full optimizer state dict in place of using ``torch.save()`` and ``torch.load()`` so that all ranks can have it.""" obj_list = [full_osd] dist.broadcast_object_list( obj_list, src=0, group=group, ) full_osd = obj_list[0] return full_osd
def _collect_sharded_states(self) -> List[Dict[str, Any]]: """Collect all the state shards, in CPU memory.""" all_states = [] for rank in range(self.world_size): if rank == self.rank: logging.debug("Saving self state") all_states.append( recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu"))) # Sync with other replicas if _torch_broadcast_object: # torch native object broadcast dist.broadcast_object_list([0], src=self.global_rank, group=self.group) else: # legacy compatibility for old torch versions broadcast_object( torch.tensor([0], dtype=torch.uint8, device=self._device), src_rank=self.global_rank, group=self.group, dist_device=self._device, ) else: # Fetch the optim state from the other replicas global_rank = self.get_global_rank(self.group, rank) if _torch_broadcast_object: replica_state_l = [0] dist.broadcast_object_list(replica_state_l, src=global_rank, group=self.group) replica_state = replica_state_l[0] else: replica_state = broadcast_object( torch.tensor([0], dtype=torch.uint8, device=self._device), src_rank=global_rank, group=self.group, dist_device=self._device, ) all_states.append( recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))) logging.debug("State from rank %s received", rank) return all_states
def run_test_collect_shards(rank, world_size, reference_rank, tempfile_name): dist_init(rank, world_size, tempfile_name) device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE # Run a dummy step so that the optimizer state dict exists batch, input_width, hidden, target_width = 3, 3, 3, 5 target = torch.rand((batch, target_width), device=device) inputs = torch.rand((batch, input_width), device=device) model = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, target_width)) model.to(device) loss_fn = torch.nn.L1Loss() loss_fn.to(device) # With SGD, Momentum is required to get a state to shard optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99) def closure(): optimizer.zero_grad() output = model(inputs) loss = loss_fn(output, target) loss.backward() return loss _ = optimizer.step(closure=closure) # Update the optimizer state on the reference rank optimizer.consolidate_state_dict(recipient_rank=reference_rank) # Fetch the state on the reference rank # - check that it has the correct size # - load it again if rank == reference_rank: optimizer_state_dict = optimizer.state_dict() assert len(optimizer_state_dict["state"]) == world_size else: optimizer_state_dict = {} optim_state = [optimizer_state_dict] if _torch_broadcast_object: dist.broadcast_object_list(optim_state, src=reference_rank, group=dist.group.WORLD) optimizer_state_dict = optim_state[0] else: optimizer_state_dict = optim.utils.broadcast_object( optimizer_state_dict, src_rank=reference_rank, group=dist.group.WORLD, dist_device=device) # Load the optimizer state dict optimizer.load_state_dict(optimizer_state_dict) dist.destroy_process_group()
def test_step(self, split="val"): if self.device == 0: metric, loss = self._test_step() val_loss = float(loss[split]) val_metric = float(metric[split]) object_list = [val_metric, val_loss] else: object_list = [None, None] dist.broadcast_object_list(object_list, src=0) return object_list[0], object_list[1]
def broadcast_object(self, object: Optional[T]) -> T: """ Same as c10d::broadcast_object_list but works without distributed enabled. """ object_list = [object] if self.use_dist: dist.broadcast_object_list(object_list=object_list, group=self.group, src=self.coordinator_rank) return cast(T, object_list[0])
def gen_metadata(self) -> Metadata: module = TestModule() # compute the default saved metadata (must pass include_non_replicated_tensors or we'll get incomplete MD) metadata, _, _ = _prepare(module.state_dict(), True) # _prepare only produc metadata = [metadata] dist.broadcast_object_list(metadata) return metadata[0]
def sync_object_ranks(something_to_sync: Any, reference_rank: int, device: torch.device) -> Any: if _torch_broadcast_object: package = [something_to_sync] dist.broadcast_object_list(package, src=reference_rank, group=dist.group.WORLD) package_sync = package[0] else: package_sync = optim.utils.broadcast_object(something_to_sync, src_rank=reference_rank, group=dist.group.WORLD, dist_device=device) return package_sync
def test_read_write_shard_tensor(self) -> None: paths = [tempfile.mkdtemp()] dist.broadcast_object_list(paths) path = paths[0] # pyre-fixme [28]: Unexpected keyword argument `dim` to call `dist._sharding_spec.api.ChunkShardingSpec.__init__`. spec = ChunkShardingSpec( dim=0, placements=[ "rank:0", "rank:1", ], ) model_to_save = MyShardedModel1(spec, init_rrefs=False) # Test save model_to_save._register_state_dict_hook(state_dict_hook) state_dict_to_save = model_to_save.state_dict() fs_writer = FileSystemWriter(path=path) save_state_dict(state_dict=state_dict_to_save, storage_writer=fs_writer) dist.barrier() # Create a new model model_to_load = MyShardedModel1(spec, init_rrefs=False) # This is not the correct hook for loading the state dict # model_to_load._register_load_state_dict_pre_hook(pre_load_state_dict_hook, True) model_to_load._register_state_dict_hook(state_dict_hook) state_dict_to_load_to = model_to_load.state_dict() dist.barrier() with self.assertRaises(AssertionError): assert_state_dict_equal(self, state_dict_to_load_to, state_dict_to_save) # Test load. fs_reader = FileSystemReader(path=path) load_state_dict(state_dict=state_dict_to_load_to, storage_reader=fs_reader) assert_state_dict_equal(self, state_dict_to_load_to, state_dict_to_save) dist.barrier()
def _broadcast_state_dict(self) -> None: """Broadcast this rank's state shard, discard others""" # Default to CPU space to gain some memory headroom local_cpu_state = recursive_copy_to_device(self.local_state_dict(), non_blocking=True, device=torch.device("cpu")) # Tensor cannot be really empty, even if its size is meaningless dummy_sync_tensor = torch.tensor([1], device=self._device) for rank in range(self.world_size): if rank == self.rank: # Send the state to the reference replica logging.debug( "Sending the sharded optimizer state to the reference replica from rank %s", rank, ) if _torch_broadcast_object: # torch native object broadcast dist.broadcast_object_list([local_cpu_state], src=self.global_rank, group=self.group) else: # legacy compatibility for old torch versions broadcast_object(self.local_state_dict(), src_rank=self.global_rank, group=self.group, dist_device=self._device) else: global_rank = self.get_global_rank(self.group, rank) # Discard this tensor/rank, broadcast necessary for syncing and because NCCL does not support gather if _torch_broadcast_object: dist.broadcast_object_list([dummy_sync_tensor], src=global_rank, group=self.group) else: broadcast_object( torch.tensor([dummy_sync_tensor], dtype=torch.uint8, device=self._device), src_rank=global_rank, group=self.group, dist_device=self._device, )
def distributed_test(self, model_w: ModelWrapper, loader, rank, fn): model_w.eval() # if rank == 0: if dist.get_rank() == 0: if self.cpu_inference: model_w.to("cpu") _device = "cpu" else: _device = rank with torch.no_grad(): result = fn(model_w, loader, _device) model_w.to(rank) object_list = [result] else: object_list = [None] dist.broadcast_object_list(object_list, src=0) return object_list[0]
def test_distributed_checkpoint(self, state_dict_type) -> None: with enable_wrap(wrapper_cls=FSDP): torch.manual_seed(100) model = wrap(SkipModel(double_nest=True)) torch.manual_seed(200) new_model = wrap(SkipModel(double_nest=True)) with FullyShardedDataParallel.summon_full_params( model), FullyShardedDataParallel.summon_full_params(new_model): params = list(model.parameters()) new_params = list(new_model.parameters()) self.assertNotEqual(params, new_params) with tempfile.TemporaryDirectory() as path: paths = [path] dist.broadcast_object_list(paths) path = paths[0] writer = FileSystemWriter(path) reader = FileSystemReader(path) with FSDP.state_dict_type(model, state_dict_type), FSDP.state_dict_type( new_model, state_dict_type): state_dict = model.state_dict() save_state_dict(state_dict, writer) with FSDP.state_dict_type(model, state_dict_type), FSDP.state_dict_type( new_model, state_dict_type): state_dict = new_model.state_dict() load_state_dict(state_dict, reader) new_model.load_state_dict(state_dict) with FullyShardedDataParallel.summon_full_params( model), FullyShardedDataParallel.summon_full_params(new_model): params = list(model.parameters()) new_params = list(new_model.parameters()) self.assertEqual(params, new_params)
def _broadcast_processed_optim_state_dict( processed_optim_state_dict: Optional[Dict[str, Any]], fsdp_flat_param_ids: Optional[Set[int]], rank: int, group, device: torch.device, ) -> Tuple[Dict[str, Any], Set[int]]: """ Broadcasts the processed optimizer state dict and the accompanying FSDP parameter IDs from rank 0 to all ranks. Args: processed_optim_state_dict (Optional[Dict[str, Any]]): The full optimizer state dict with positive-dimension tensor states replaced with metadata if on rank 0; ignored otherwise. fsdp_flat_param_ids (Optional[Set[int]]): Parameter IDs corresponding to FSDP parameters if on rank 0; ignored otherwise. device (torch.device): Device to move zero-dimension tensors post- broadcast. Returns: Tuple[Dict[str, Any], Set[int]]: The processed optimizer state dict and the parameter IDs corresponding to FSDP parameters. """ # Broadcast the two data structures rank 0 to all ranks obj_list = [processed_optim_state_dict, fsdp_flat_param_ids] if rank == 0 \ else [None, None] dist.broadcast_object_list(obj_list, src=0, group=group) processed_optim_state_dict, fsdp_flat_param_ids = obj_list # type: ignore[assignment] assert processed_optim_state_dict is not None assert fsdp_flat_param_ids is not None # Move zero-dimension tensors to `device` for param_state in processed_optim_state_dict["state"].values(): for state_name, value in param_state.items(): if _is_zero_dim_tensor(value): param_state[state_name] = value.to(device) return processed_optim_state_dict, fsdp_flat_param_ids
def consolidate_state_dict(self, recipient_rank: int = 0) -> None: """Update the consolidated state_dict list, one per rank. .. warning: This needs to be called on all replicas""" # Sync lr and other attributes in case its been updated self._update_param_groups() empty_messenger = torch.tensor([0], dtype=torch.uint8, device=self._device) # Pull the sharded state from all the other replicas # Store all the states in order, rank by rank # NOTE: In practice, `broadcast` is used, which is wasteful (gather would have been appropriate) # compatibility issues with some backends make the use of broadcast mandatory for now. # a possible follow up would be to move all sharded state management to RPC RRef self._all_states = [] for rank in range(self.world_size): global_rank = _get_global_rank(self.group, rank) # This rank collects the whole state if self.rank == recipient_rank: if rank == self.rank: self._all_states.append( _recursive_copy_to_device(self.optim.state_dict(), non_blocking=True, device=torch.device("cpu"))) else: # Fetch the optim state from the other replicas replica_state = [empty_messenger] dist.broadcast_object_list(object_list=replica_state, src=global_rank, group=self.group) self._all_states.append( _recursive_copy_to_device(replica_state[0], non_blocking=True, device=torch.device("cpu"))) else: # Acknowledge broadcasts, and send this rank's shard when needed # Default to CPU space to gain some memory headroom if rank == self.rank: # Send the state to the reference replica dist.broadcast_object_list( object_list=[self.optim.state_dict()], src=self.global_rank, group=self.group) elif rank != recipient_rank: # Discard this tensor/rank, broadcast was being use for compatibility reasons dist.broadcast_object_list(object_list=[empty_messenger], src=global_rank, group=self.group)
def _broadcast_processed_optim_state_dict( processed_optim_state_dict: Optional[Dict[str, Any]], rank: int, group, ) -> Dict[str, Any]: """ Broadcasts the processed optimizer state dict from rank 0 to all ranks. Args: processed_optim_state_dict (Optional[Dict[str, Any]]): The flattened optimizer state dict with positive-dimension tensor states replaced with metadata if on rank 0; ignored otherwise. Returns: Dict[str, Any]: The processed optimizer state dict. """ # Broadcast the two data structures rank 0 to all ranks obj_list = [processed_optim_state_dict] if rank == 0 \ else [None] dist.broadcast_object_list(obj_list, src=0, group=group) processed_optim_state_dict = obj_list[0] # type: ignore[assignment] assert processed_optim_state_dict is not None # Keep zero-dimension tensors on CPU return processed_optim_state_dict
def _broadcast_list_from_main(self, val: Any) -> Any: if not self._is_ddp: return val dist.broadcast_object_list(val, src=0, group=self._gloo_handle) return val
def _broadcast_state_dict(self, state_dict): olist = [state_dict if self.rank == 0 else None] dist.broadcast_object_list(olist) return olist[0]
def check_optimizer_equivalence( optimizer: Type[torch.optim.Optimizer]): # Any model works. Add one different buffer per rank model = torch.nn.Sequential( torch.nn.Linear(2, 3), torch.nn.Linear(3, 3), torch.nn.Linear(3, 3), ) model.register_buffer("test_buffer", torch.ones((1)) * self.rank) model.to(self.device) sharded_optimizer = ZeroRedundancyOptimizer( params=model.parameters(), optimizer_class=optimizer, lr=1e-3) sharded_ddp_model = DDP(module=model, device_ids=[self.rank], broadcast_buffers=True, find_unused_parameters=True) ddp_model_single = copy.deepcopy(model) ddp_model_single.to(self.device) ddp_optimizer = optimizer(ddp_model_single.parameters(), lr=1e-3) ddp_model = DDP(ddp_model_single, device_ids=[self.rank], broadcast_buffers=True, find_unused_parameters=True) # The model should be synchronized in between the ranks at construction time, check that check_same_model_params(sharded_ddp_model, ddp_model, "Models differ from the start") def check_step(): input_tensor = torch.rand((64, 2)) def closure_ddp(input_tensor=input_tensor): ddp_optimizer.zero_grad() ddp_loss = ddp_model(input_tensor).abs().sum() ddp_loss.backward() return ddp_loss def closure_sharded(input_tensor=input_tensor): sharded_optimizer.zero_grad() sharded_loss = sharded_ddp_model( input_tensor).abs().sum() sharded_loss.backward() return sharded_loss loss_ddp = cast(torch.Tensor, ddp_optimizer.step(closure=closure_ddp)) loss_sharded_optim = cast( torch.Tensor, sharded_optimizer.step(closure=closure_sharded)) assert torch.allclose( loss_ddp, loss_sharded_optim ), "Losses differ in between Pytorch optim and ZeroRedundancyOptimizer" check_same_model_params(sharded_ddp_model, ddp_model, "Models differ after a step") # The models should stay the same in between the ranks for i in range(BATCHS): check_step() # Change the models trainability, check that parity is maintained # only check after a couple of constant batchs to go through both regimes if i > BATCHS // 2: next(ddp_model.parameters()).requires_grad = bool(i % 2) next(sharded_ddp_model.parameters() ).requires_grad = bool(i % 2) # Check that the checkpoints are compatible reference_rank = 0 # - get states ddp_state_dict = ddp_optimizer.state_dict() sharded_optimizer.consolidate_state_dict(to=reference_rank) sharded_optim_state_dict = [ sharded_optimizer.state_dict() if self.rank == reference_rank else {} ] dist.broadcast_object_list(sharded_optim_state_dict, src=reference_rank, group=dist.group.WORLD) sharded_optim_state_dict = sharded_optim_state_dict[0] # - cross load the states # run one step and check that the models are still the same ddp_state_dict_ref = copy.deepcopy( ddp_state_dict) # OSS will remove some states ddp_optimizer.load_state_dict( sharded_optim_state_dict) # mixup on purpose ! sharded_optimizer.load_state_dict(ddp_state_dict) check_step() # - self load, rewind, check no problem # run one step and check that the models are still the same ddp_optimizer.load_state_dict(ddp_state_dict_ref) sharded_optimizer.load_state_dict(sharded_optim_state_dict) check_step()
import torch.distributed as dist if dist.get_rank() == 0: objects = ["f", 1] else: objects = [None, None] # ruleid: pickles-in-torch-distributed dist.broadcast_object_list(objects, src=0) # ruleid: pickles-in-torch-distributed dist.all_gather_object(output, gather_objects[dist.get_rank()]) # ruleid: pickles-in-torch-distributed dist.gather_object(gather_objects[dist.get_rank()], output if dist.get_rank() == 0 else None, dst=0) # ruleid: pickles-in-torch-distributed dist.scatter_object_list(output_list, objects, src=0)
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \ Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \ opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze callbacks.run('on_pretrain_routine_start') # Directories w = save_dir / 'weights' # weights dir (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir last, best = w / 'last.pt', w / 'best.pt' # Hyperparameters if isinstance(hyp, str): with open(hyp, errors='ignore') as f: hyp = yaml.safe_load(f) # load hyps dict LOGGER.info( colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) # Save run settings if not evolve: with open(save_dir / 'hyp.yaml', 'w') as f: yaml.safe_dump(hyp, f, sort_keys=False) with open(save_dir / 'opt.yaml', 'w') as f: yaml.safe_dump(vars(opt), f, sort_keys=False) # Loggers data_dict = None if RANK in {-1, 0}: loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance if loggers.wandb: data_dict = loggers.wandb.data_dict if resume: weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size # Register actions for k in methods(loggers): callbacks.register_action(k, callback=getattr(loggers, k)) # Config plots = not evolve and not opt.noplots # create plots cuda = device.type != 'cpu' init_seeds(opt.seed + 1 + RANK, deterministic=True) with torch_distributed_zero_first(LOCAL_RANK): data_dict = data_dict or check_dataset(data) # check if None train_path, val_path = data_dict['train'], data_dict['val'] nc = 1 if single_cls else int(data_dict['nc']) # number of classes names = ['item'] if single_cls and len( data_dict['names']) != 1 else data_dict['names'] # class names assert len( names ) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check is_coco = isinstance(val_path, str) and val_path.endswith( 'coco/val2017.txt') # COCO dataset # Model check_suffix(weights, '.pt') # check weights pretrained = weights.endswith('.pt') if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download( weights) # download if not found locally ckpt = torch.load(weights, map_location='cpu' ) # load checkpoint to CPU to avoid CUDA memory leak model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create exclude = [ 'anchor' ] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys csd = ckpt['model'].float().state_dict( ) # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load LOGGER.info( f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}' ) # report else: model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create amp = check_amp(model) # check AMP # Freeze freeze = [ f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0])) ] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers if any(x in k for x in freeze): LOGGER.info(f'freezing {k}') v.requires_grad = False # Image size gs = max(int(model.stride.max()), 32) # grid size (max stride) imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple # Batch size if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size batch_size = check_train_batch_size(model, imgsz, amp) loggers.on_params_update({"batch_size": batch_size}) # Optimizer nbs = 64 # nominal batch size accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}") optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay']) # Scheduler if opt.cos_lr: lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] else: lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf' ] # linear scheduler = lr_scheduler.LambdaLR( optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) # EMA ema = ModelEMA(model) if RANK in {-1, 0} else None # Resume start_epoch, best_fitness = 0, 0.0 if pretrained: # Optimizer if ckpt['optimizer'] is not None: optimizer.load_state_dict(ckpt['optimizer']) best_fitness = ckpt['best_fitness'] # EMA if ema and ckpt.get('ema'): ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) ema.updates = ckpt['updates'] # Epochs start_epoch = ckpt['epoch'] + 1 if resume: assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.' if epochs < start_epoch: LOGGER.info( f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs." ) epochs += ckpt['epoch'] # finetune additional epochs del ckpt, csd # DP mode if cuda and RANK == -1 and torch.cuda.device_count() > 1: LOGGER.warning( 'WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n' 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.' ) model = torch.nn.DataParallel(model) # SyncBatchNorm if opt.sync_bn and cuda and RANK != -1: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) LOGGER.info('Using SyncBatchNorm()') # Trainloader train_loader, dataset = create_dataloader( train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls, hyp=hyp, augment=True, cache=None if opt.cache == 'val' else opt.cache, rect=opt.rect, rank=LOCAL_RANK, workers=workers, image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '), shuffle=True) mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class nb = len(train_loader) # number of batches assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}' # Process 0 if RANK in {-1, 0}: val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls, hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1, workers=workers * 2, pad=0.5, prefix=colorstr('val: '))[0] if not resume: labels = np.concatenate(dataset.labels, 0) # c = torch.tensor(labels[:, 0]) # classes # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency # model._initialize_biases(cf.to(device)) if plots: plot_labels(labels, names, save_dir) # Anchors if not opt.noautoanchor: check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) model.half().float() # pre-reduce anchor precision callbacks.run('on_pretrain_routine_end') # DDP mode if cuda and RANK != -1: model = smart_DDP(model) # Model attributes nl = de_parallel( model).model[-1].nl # number of detection layers (to scale hyps) hyp['box'] *= 3 / nl # scale to layers hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers hyp['obj'] *= (imgsz / 640)**2 * 3 / nl # scale to image size and layers hyp['label_smoothing'] = opt.label_smoothing model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model model.class_weights = labels_to_class_weights( dataset.labels, nc).to(device) * nc # attach class weights model.names = names # Start training t0 = time.time() nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations) # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training last_opt_step = -1 maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0 ) # P, R, [email protected], [email protected], val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move scaler = torch.cuda.amp.GradScaler(enabled=amp) stopper, stop = EarlyStopping(patience=opt.patience), False compute_loss = ComputeLoss(model) # init loss class callbacks.run('on_train_start') LOGGER.info( f'Image sizes {imgsz} train, {imgsz} val\n' f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n' f"Logging results to {colorstr('bold', save_dir)}\n" f'Starting training for {epochs} epochs...') for epoch in range( start_epoch, epochs ): # epoch ------------------------------------------------------------------ callbacks.run('on_train_epoch_start') model.train() # Update image weights (optional, single-GPU only) if opt.image_weights: cw = model.class_weights.cpu().numpy() * ( 1 - maps)**2 / nc # class weights iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx # Update mosaic border (optional) # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) # dataset.mosaic_border = [b - imgsz, -b] # height, width borders mloss = torch.zeros(3, device=device) # mean losses if RANK != -1: train_loader.sampler.set_epoch(epoch) pbar = enumerate(train_loader) LOGGER.info( ('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size')) if RANK in {-1, 0}: pbar = tqdm( pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar optimizer.zero_grad() for i, ( imgs, targets, paths, _ ) in pbar: # batch ------------------------------------------------------------- callbacks.run('on_train_batch_start') ni = i + nb * epoch # number integrated batches (since train start) imgs = imgs.to(device, non_blocking=True).float( ) / 255 # uint8 to float32, 0-255 to 0.0-1.0 # Warmup if ni <= nw: xi = [0, nw] # x interp # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) accumulate = max( 1, np.interp(ni, xi, [1, nbs / batch_size]).round()) for j, x in enumerate(optimizer.param_groups): # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 x['lr'] = np.interp(ni, xi, [ hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch) ]) if 'momentum' in x: x['momentum'] = np.interp( ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) # Multi-scale if opt.multi_scale: sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size sf = sz / max(imgs.shape[2:]) # scale factor if sf != 1: ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:] ] # new shape (stretched to gs-multiple) imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) # Forward with torch.cuda.amp.autocast(amp): pred = model(imgs) # forward loss, loss_items = compute_loss( pred, targets.to(device)) # loss scaled by batch_size if RANK != -1: loss *= WORLD_SIZE # gradient averaged between devices in DDP mode if opt.quad: loss *= 4. # Backward scaler.scale(loss).backward() # Optimize if ni - last_opt_step >= accumulate: scaler.step(optimizer) # optimizer.step scaler.update() optimizer.zero_grad() if ema: ema.update(model) last_opt_step = ni # Log if RANK in {-1, 0}: mloss = (mloss * i + loss_items) / (i + 1 ) # update mean losses mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB) pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % (f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])) callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots) if callbacks.stop_training: return # end batch ------------------------------------------------------------------------------------------------ # Scheduler lr = [x['lr'] for x in optimizer.param_groups] # for loggers scheduler.step() if RANK in {-1, 0}: # mAP callbacks.run('on_train_epoch_end', epoch=epoch) ema.update_attr(model, include=[ 'yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights' ]) final_epoch = (epoch + 1 == epochs) or stopper.possible_stop if not noval or final_epoch: # Calculate mAP results, maps, _ = val.run(data_dict, batch_size=batch_size // WORLD_SIZE * 2, imgsz=imgsz, model=ema.ema, single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, plots=False, callbacks=callbacks, compute_loss=compute_loss) # Update best mAP fi = fitness(np.array(results).reshape( 1, -1)) # weighted combination of [P, R, [email protected], [email protected]] stop = stopper(epoch=epoch, fitness=fi) # early stop check if fi > best_fitness: best_fitness = fi log_vals = list(mloss) + list(results) + lr callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi) # Save model if (not nosave) or (final_epoch and not evolve): # if save ckpt = { 'epoch': epoch, 'best_fitness': best_fitness, 'model': deepcopy(de_parallel(model)).half(), 'ema': deepcopy(ema.ema).half(), 'updates': ema.updates, 'optimizer': optimizer.state_dict(), 'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None, 'date': datetime.now().isoformat() } # Save last, best and delete torch.save(ckpt, last) if best_fitness == fi: torch.save(ckpt, best) if opt.save_period > 0 and epoch % opt.save_period == 0: torch.save(ckpt, w / f'epoch{epoch}.pt') del ckpt callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi) # EarlyStopping if RANK != -1: # if DDP training broadcast_list = [stop if RANK == 0 else None] dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks if RANK != 0: stop = broadcast_list[0] if stop: break # must break all DDP ranks # end epoch ---------------------------------------------------------------------------------------------------- # end training ----------------------------------------------------------------------------------------------------- if RANK in {-1, 0}: LOGGER.info( f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.' ) for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers if f is best: LOGGER.info(f'\nValidating {f}...') results, _, _ = val.run( data_dict, batch_size=batch_size // WORLD_SIZE * 2, imgsz=imgsz, model=attempt_load(f, device).half(), iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65 single_cls=single_cls, dataloader=val_loader, save_dir=save_dir, save_json=is_coco, verbose=True, plots=plots, callbacks=callbacks, compute_loss=compute_loss) # val best model with plots if is_coco: callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi) callbacks.run('on_train_end', last, best, plots, epoch, results) torch.cuda.empty_cache() return results
def train(model, dataset, args, rank, evaluator, loss_fn): print(f"Running on rank {rank}.") os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(args.master_port) # initialize the process group dist.init_process_group("nccl", rank=rank, world_size=args.world_size) model = copy.deepcopy(model).to(rank) model = DDP(model, device_ids=[rank]) data = dataset[0] train_dataset, train_loader = get_train_loader(dataset, args, rank) optimizer = torch.optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) epoch_iter = tqdm(range(args.max_epoch)) if rank == 0 else range( args.max_epoch) patience = 0 max_score = 0 min_loss = np.inf best_model = None for epoch in epoch_iter: train_dataset.shuffle() train_step(model, train_loader, optimizer, rank) if (epoch + 1) % args.eval_step == 0: if rank == 0: acc, loss = test_step(model.module, data, evaluator, loss_fn) train_acc = acc["train"] val_acc = acc["val"] val_loss = loss["val"] epoch_iter.set_description( f"Epoch: {epoch:03d}, Train: {train_acc:.4f}, Val: {val_acc:.4f}" ) model = model.to(rank) object_list = [val_loss, val_acc] else: object_list = [None, None] dist.broadcast_object_list(object_list, src=0) val_loss, val_acc = object_list if val_loss <= min_loss or val_acc >= max_score: if val_loss <= min_loss: best_model = copy.deepcopy(model) min_loss = np.min((min_loss, val_loss)) max_score = np.max((max_score, val_acc)) patience = 0 else: patience += 1 if patience == args.patience: break dist.barrier() if rank == 0: os.makedirs("./checkpoints", exist_ok=True) checkpoint_path = os.path.join("./checkpoints", f"{args.model}_{args.dataset}.pt") if best_model is not None: print(f"Saving model to {checkpoint_path}") torch.save(best_model.module.state_dict(), checkpoint_path) dist.barrier() dist.destroy_process_group()
def consolidate_state_dict(self, recipient_rank: int = 0) -> None: """Update the consolidated state_dict list, one per rank. Arguments: recipient_rank (int): on which rank to materialize the full state dict. -1 is a special value, which means that all ranks should have the state .. warning: This needs to be called on all replicas""" # Sync lr and other attributes in case its been updated OSS._sync_param_groups(self.param_groups, self.optim.param_groups) # Pull the sharded state from all the other replicas # Store all the states in order, rank by rank logging.debug("Pulling the sharded optimizer state from all replicas") self._all_states = [] should_collect_state = self.rank == recipient_rank or recipient_rank == -1 should_send_state = self.rank != recipient_rank # NCCL requires CUDA tensors for all communication primitives dist_device = torch.device( "cuda" ) if self.backend == dist.Backend.NCCL else self._default_device for rank in range(self.world_size): if rank == self.rank: if should_collect_state: logging.debug("Saving self state") self._all_states.append( recursive_copy_to_device(self.optim.state_dict(), non_blocking=True, device=torch.device("cpu"))) # Sync with other replicas state_to_share = ( self.optim.state_dict() if should_send_state else torch.tensor([0], dtype=torch.uint8, device=dist_device)) if _gpu_capabilities_older_than_50(): _broadcast_object(state_to_share, src_rank=self.global_rank, group=self.group, dist_device=dist_device) else: obj_list = [state_to_share] dist.broadcast_object_list( obj_list, src=self.global_rank, group=self.group, ) else: # Fetch the optim state from the other replicas if _gpu_capabilities_older_than_50(): replica_state = _broadcast_object( torch.tensor([0], dtype=torch.uint8, device=dist_device), src_rank=self._local_to_global_rank[rank], group=self.group, dist_device=dist_device, ) else: obj_list = [ torch.tensor([0], dtype=torch.uint8, device=dist_device) ] dist.broadcast_object_list( obj_list, src=self._local_to_global_rank[rank], group=self.group, ) replica_state = obj_list[0] if should_collect_state: self._all_states.append( recursive_copy_to_device(replica_state, non_blocking=True, device=torch.device("cpu"))) logging.debug("State from rank %s received", rank)
def save_state_dict( state_dict: Dict[str, Any], storage_writer: StorageWriter, process_group: Optional[dist.ProcessGroup] = None, coordinator_rank: int = 0, no_dist: bool = False ) -> None: """ Save a distributed model in SPMD style. This function is different from ``torch.save()`` as it handles ``ShardedTensor`` by having each rank only save their local shards. To produce a state_dict with ShardedTensor instances you must call ``_register_state_dict_hook`` on the top module with value `torch.distributed._shard.sharded_tensor.state_dict_hook` prior to calling `state_dict()` on the top module. There is no guarantees of Backwards Compatibility across PyTorch versions for saved state_dicts. If using the `process_group` argument, make sure that only its ranks call `save_state_dict` and that all data in state_dict belong to it. This function can be used to save a state_dict with an intialized process group by passing ``no_dist=True``. This can be used to produce a checkpoint that can consumed by load_state_dict is a SPMD fashion. Args: state_dict (Dict[str, Any]) : A state_dict storage_writer (StorageWriter): Instance of StorageWrite use to perform writes. process_group (ProcessGroup): ProcessGroup to be used for cross-rank synchronization coordinator_rank (int): Rank to use to coordinate the checkpoint, rank0 is used by default no_dist (bool): Don't attempt to save in SPMD style. Default to False Example: >>> my_model = MyModule() >>> # We must call this function prior to state_dict() >>> my_model._register_state_dict_hook(state_dict_hook) >>> model_state_dict = my_model.state_dict() >>> fs_storage_writer = torch.distributed._shard.checkpoint.FileSystemWriter("/checkpoint/1") >>> torch.distributed._shard.checkpoint.save_state_dict( >>> state_dict=model_state_dict, >>> storage_writer=fs_stroage_writer, >>> ) .. note:: save_state_dict uses collectives to coordinate writes across ranks. For NCCL-based process groups, internal tensor representations of objects must be moved to the GPU device before communication takes place. In this case, the device used is given by ``torch.cuda.current_device()`` and it is the user's responsibility to ensure that this is set so that each rank has an individual GPU, via ``torch.cuda.set_device()`` """ is_coordinator = no_dist or dist.get_rank(process_group) == coordinator_rank exceptions: List[Optional[BaseException]] = [None] if is_coordinator: try: storage_writer.prepare() except BaseException as e: exceptions = [e] # Writing can only start once prepare has finished if not no_dist: dist.broadcast_object_list(exceptions, group=process_group, src=coordinator_rank) if exceptions[0] is not None: raise CheckpointException("failed to prepare storage", {coordinator_rank : exceptions[0]}) rank_write_error: Optional[BaseException] try: ( metadata, bytes_write_requests, tensor_write_requests, ) = _prepare(state_dict, is_coordinator, process_group) combined_writes: List[Union[TensorWriteRequest, BytesWriteRequest]] = [] combined_writes.extend(tensor_write_requests) combined_writes.extend(bytes_write_requests) storage_writer.prepare_storage(combined_writes) bytes_futures = storage_writer.write_bytes(bytes_write_requests) tensor_futures = storage_writer.write_tensors(tensor_write_requests) torch.futures.wait_all([bytes_futures, tensor_futures]) rank_write_error = None except BaseException as e: rank_write_error = e all_errors: List[Optional[BaseException]] # collect all write errors if not no_dist: all_errors = [None] * dist.get_world_size(process_group) dist.gather_object( obj=rank_write_error, object_gather_list=all_errors if is_coordinator else None, dst=coordinator_rank ) else: all_errors = [rank_write_error] result: List[Optional[CheckpointException]] = [None] if is_coordinator: message: Optional[str] = None # gather produces an array of arrays, flatten it if any(all_errors): message = "Failed to write data" else: try: storage_writer.finish(metadata=metadata) except BaseException as e: all_errors[coordinator_rank] = e message = "Failed to finish checkpoint" if message is not None: node_failures = {i: err for i, err in enumerate(all_errors) if err is not None} result[0] = CheckpointException(message, node_failures) if not no_dist: dist.broadcast_object_list( result, group=process_group, src=coordinator_rank) if result[0] is not None: raise result[0]
def check_optimizer_equivalence( optimizer: Type[torch.optim.Optimizer]): # Any model works. Add one different buffer per rank model = torch.nn.Sequential( torch.nn.Linear(2, 3), torch.nn.Linear(3, 3), torch.nn.Linear(3, 3), ) model.register_buffer("test_buffer", torch.ones((1)) * self.rank) model.to(self.device) sharded_optimizer = ZeroRedundancyOptimizer( params=model.parameters(), optim=optimizer, lr=1e-3) sharded_ddp_model = DDP(module=model, device_ids=[self.rank], broadcast_buffers=True) ddp_model_single = copy.deepcopy(model) ddp_model_single.to(self.device) ddp_optimizer = optimizer(ddp_model_single.parameters(), lr=1e-3) ddp_model = DDP(ddp_model_single, device_ids=[self.rank], broadcast_buffers=True) def check_same_model_params(): for pg, ddp_pg in zip(sharded_optimizer.param_groups, ddp_optimizer.param_groups): for p, ddp_p in zip(pg["params"], ddp_pg["params"]): assert torch.allclose( p, ddp_p, atol=1e-3 ), f"Model parameters differ in between Pytorch optim and ZeroRedundancyOptimizer \n{p} {ddp_p}" for b, ddp_b in zip(sharded_ddp_model.buffers(), ddp_model.buffers()): assert torch.allclose( b, ddp_b ), "Model buffers differ in between Pytorch optim and ZeroRedundancyOptimizer" # The model should be synchronized in between the ranks at construction time, check that check_same_model_params() # The models should stay the same across multiple steps, losses should stay the same def check_step(): input_tensor = torch.rand((64, 2)) def closure_ddp(input_tensor=input_tensor): ddp_optimizer.zero_grad() ddp_loss = ddp_model(input_tensor).abs().sum() ddp_loss.backward() return ddp_loss def closure_sharded(input_tensor=input_tensor): sharded_optimizer.zero_grad() sharded_loss = sharded_ddp_model( input_tensor).abs().sum() sharded_loss.backward() return sharded_loss loss_ddp = cast(torch.Tensor, ddp_optimizer.step(closure=closure_ddp)) loss_sharded_optim = cast( torch.Tensor, sharded_optimizer.step(closure=closure_sharded)) assert torch.allclose( loss_ddp, loss_sharded_optim ), "Losses differ in between Pytorch optim and ZeroRedundancyOptimizer" check_same_model_params() for i in range(20): check_step() # Test state dict save/load/equivalence with pytorch # - save state for both sharded_optimizer.consolidate_state_dict() sharded_optimizer_state_dict = (sharded_optimizer.state_dict() if self.rank == RECIPIENT_RANK else torch.zeros(1)) ddp_state_dict = ddp_optimizer.state_dict() # - sync the saved state with all the ranks exchange_list = [sharded_optimizer_state_dict] dist.broadcast_object_list( exchange_list, src=RECIPIENT_RANK, group=dist.group.WORLD, ) sharded_optimizer_state_dict = exchange_list[0] # - cross load the states ddp_optimizer.load_state_dict(sharded_optimizer_state_dict) sharded_optimizer.load_state_dict(ddp_state_dict) # - run one step, and check that the models are still the same check_step()
def get_file_path(self) -> str: paths = [tempfile.mkdtemp()] if dist.get_rank() == 0 else [None] dist.broadcast_object_list(paths) return paths[0]
def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer]): # Any model works. Add one different buffer per rank model = torch.nn.Sequential( torch.nn.Linear(2, 3), torch.nn.Linear(3, 3), torch.nn.Linear(3, 3), ) model.register_buffer("test_buffer", torch.ones((1)) * self.rank) model.to(self.device) sharded_optimizer = ZeroRedundancyOptimizer(params=model.parameters(), optim=optimizer, lr=1e-3) sharded_ddp_model = DDP(module=model, device_ids=[self.rank], broadcast_buffers=True) ddp_model_single = copy.deepcopy(model) ddp_model_single.to(self.device) ddp_optimizer = optimizer(ddp_model_single.parameters(), lr=1e-3) ddp_model = DDP(ddp_model_single, device_ids=[self.rank], broadcast_buffers=True) def check_same_model_params(): for pg, ddp_pg in zip(sharded_optimizer.param_groups, ddp_optimizer.param_groups): for p, ddp_p in zip(pg["params"], ddp_pg["params"]): assert torch.allclose( p, ddp_p, atol=1e-3 ), f"Model parameters differ in between Pytorch optim and ZeroRedundancyOptimizer \n{p} {ddp_p}" for b, ddp_b in zip(sharded_ddp_model.buffers(), ddp_model.buffers()): assert torch.allclose( b, ddp_b ), "Model buffers differ in between Pytorch optim and ZeroRedundancyOptimizer" # The model should be synchronized in between the ranks at construction time, check that check_same_model_params() def check_step(): input_tensor = torch.rand((64, 2)) def closure_ddp(input_tensor=input_tensor): ddp_optimizer.zero_grad() ddp_loss = ddp_model(input_tensor).abs().sum() ddp_loss.backward() return ddp_loss def closure_sharded(input_tensor=input_tensor): sharded_optimizer.zero_grad() sharded_loss = sharded_ddp_model(input_tensor).abs().sum() sharded_loss.backward() return sharded_loss loss_ddp = cast(torch.Tensor, ddp_optimizer.step(closure=closure_ddp)) loss_sharded_optim = cast(torch.Tensor, sharded_optimizer.step(closure=closure_sharded)) assert torch.allclose( loss_ddp, loss_sharded_optim ), "Losses differ in between Pytorch optim and ZeroRedundancyOptimizer" check_same_model_params() # The models should stay the same in between the ranks for i in range(20): check_step() # Check that the checkpoints are compatible reference_rank = 0 # - get states ddp_state_dict = ddp_optimizer.state_dict() sharded_optimizer.consolidate_state_dict(recipient_rank=reference_rank) sharded_optim_state_dict = [sharded_optimizer.state_dict() if self.rank == reference_rank else {}] dist.broadcast_object_list(sharded_optim_state_dict, src=reference_rank, group=dist.group.WORLD) sharded_optim_state_dict = sharded_optim_state_dict[0] # - cross load the states # run one step and check that the models are still the same ddp_state_dict_ref = copy.deepcopy(ddp_state_dict) # OSS will remove some states ddp_optimizer.load_state_dict(sharded_optim_state_dict) # mixup on purpose ! sharded_optimizer.load_state_dict(ddp_state_dict) check_step() # - self load, rewind, check no problem # run one step and check that the models are still the same ddp_optimizer.load_state_dict(ddp_state_dict_ref) sharded_optimizer.load_state_dict(sharded_optim_state_dict) check_step()