def test_deepspeed_multigpu_stage_3_resume_training(tmpdir): """Test to ensure with Stage 3 and single GPU that we can resume training.""" initial_model = ModelParallelClassificationModel() dm = ClassifDataModule() ck = ModelCheckpoint(monitor="val_acc", mode="max", save_last=True, save_top_k=-1) initial_trainer = Trainer( default_root_dir=tmpdir, max_epochs=1, limit_train_batches=2, limit_val_batches=2, limit_test_batches=2, strategy=DeepSpeedStrategy(stage=3), gpus=1, precision=16, callbacks=[ck], enable_progress_bar=False, enable_model_summary=False, ) initial_trainer.fit(initial_model, datamodule=dm) class TestCallback(Callback): def on_train_batch_start( self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int ) -> None: original_deepspeed_strategy = initial_trainer.strategy current_deepspeed_strategy = trainer.strategy assert isinstance(original_deepspeed_strategy, DeepSpeedStrategy) assert isinstance(current_deepspeed_strategy, DeepSpeedStrategy) # assert optimizer states are the correctly loaded original_optimizer_dict = original_deepspeed_strategy.deepspeed_engine.optimizer.state_dict() current_optimizer_dict = current_deepspeed_strategy.deepspeed_engine.optimizer.state_dict() for orig_tensor, current_tensor in zip( original_optimizer_dict["fp32_flat_groups"], current_optimizer_dict["fp32_flat_groups"] ): assert torch.all(orig_tensor.eq(current_tensor)) # assert model state is loaded correctly for current_param, initial_param in zip(pl_module.parameters(), initial_model.parameters()): assert torch.equal(current_param.cpu(), initial_param.cpu()) # assert epoch has correctly been restored assert trainer.current_epoch == 1 # assert lr-scheduler states are loaded correctly original_lr_scheduler = initial_trainer.lr_scheduler_configs[0].scheduler current_lr_scheduler = trainer.lr_scheduler_configs[0].scheduler assert original_lr_scheduler.state_dict() == current_lr_scheduler.state_dict() model = ModelParallelClassificationModel() trainer = Trainer( default_root_dir=tmpdir, fast_dev_run=True, strategy=DeepSpeedStrategy(stage=3), gpus=1, precision=16, callbacks=TestCallback(), enable_progress_bar=False, enable_model_summary=False, ) trainer.fit(model, datamodule=dm, ckpt_path=ck.best_model_path)
def test_deepspeed_multigpu_single_file(tmpdir): """Test to ensure that DeepSpeed loads from a single file checkpoint.""" model = BoringModel() checkpoint_path = os.path.join(tmpdir, "model.pt") trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True) trainer.fit(model) trainer.save_checkpoint(checkpoint_path) trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=1, fast_dev_run=True, precision=16, ) strategy = trainer.strategy assert isinstance(strategy, DeepSpeedStrategy) assert not strategy.load_full_weights with pytest.raises(MisconfigurationException, match="DeepSpeed was unable to load the checkpoint."): trainer.test(model, ckpt_path=checkpoint_path) trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3, load_full_weights=True), accelerator="gpu", devices=1, fast_dev_run=True, precision=16, ) strategy = trainer.strategy assert isinstance(strategy, DeepSpeedStrategy) assert strategy.load_full_weights trainer.test(model, ckpt_path=checkpoint_path)
def test_deepspeed_multigpu_stage_3_checkpointing(tmpdir, automatic_optimization, accumulate_grad_batches): seed_everything(1) if automatic_optimization: model = ModelParallelClassificationModel() else: model = ManualModelParallelClassificationModel() dm = ClassifDataModule() ck = ModelCheckpoint(monitor="val_acc", mode="max", save_last=True, save_top_k=-1) trainer = Trainer( default_root_dir=tmpdir, max_epochs=10, strategy=DeepSpeedStrategy(stage=3), gpus=2, precision=16, accumulate_grad_batches=accumulate_grad_batches, callbacks=[ck], ) trainer.fit(model, datamodule=dm) results = trainer.test(datamodule=dm) assert results[0]["test_acc"] > 0.7 saved_results = trainer.test(ckpt_path=ck.best_model_path, datamodule=dm) assert saved_results[0]["test_acc"] > 0.7 assert saved_results == results if automatic_optimization: model = ModelParallelClassificationModel() else: model = ManualModelParallelClassificationModel() trainer = Trainer(default_root_dir=tmpdir, gpus=2, strategy=DeepSpeedStrategy(stage=3), precision=16) results = trainer.test(model, datamodule=dm, ckpt_path=ck.best_model_path) assert results[0]["test_acc"] > 0.7
def test_deepspeed_run_configure_optimizers(tmpdir): """Test end to end that deepspeed works with defaults (without ZeRO as that requires compilation), whilst using configure_optimizers for optimizers and schedulers.""" class TestCB(Callback): def on_train_start(self, trainer, pl_module) -> None: from deepspeed.runtime.zero.stage2 import FP16_DeepSpeedZeroOptimizer assert isinstance(trainer.optimizers[0], FP16_DeepSpeedZeroOptimizer) assert isinstance(trainer.optimizers[0].optimizer, torch.optim.SGD) assert isinstance(trainer.lr_scheduler_configs[0].scheduler, torch.optim.lr_scheduler.StepLR) # check that the lr_scheduler config was preserved assert trainer.lr_scheduler_configs[0].name == "Sean" class TestModel(BoringModel): def configure_optimizers(self): [optimizer], [scheduler] = super().configure_optimizers() return {"optimizer": optimizer, "lr_scheduler": {"scheduler": scheduler, "name": "Sean"}} model = TestModel() lr_monitor = LearningRateMonitor() trainer = Trainer( strategy=DeepSpeedStrategy(), # disable ZeRO so our optimizers are not wrapped default_root_dir=tmpdir, gpus=1, fast_dev_run=True, precision=16, callbacks=[TestCB(), lr_monitor], ) trainer.fit(model) assert lr_monitor.lrs == {"Sean": [0.1]} _assert_save_model_is_equal(model, tmpdir, trainer)
def test_deepspeed_auto_batch_size_config_select(mock_deepspeed_distributed, tmpdir, dataset_cls, value): """Test to ensure that the batch size is correctly set as expected for deepspeed logging purposes.""" class TestModel(BoringModel): def train_dataloader(self): return DataLoader(dataset_cls(32, 64)) class AssertCallback(Callback): def setup(self, trainer, pl_module, stage: Optional[str] = None) -> None: assert isinstance(trainer.strategy, DeepSpeedStrategy) config = trainer.strategy.config # int value overrides auto mode expected_value = value if isinstance(value, int) else 1 if dataset_cls == RandomDataset: expected_value = pl_module.train_dataloader().batch_size if value == "auto" else value assert config["train_micro_batch_size_per_gpu"] == expected_value raise SystemExit ck = AssertCallback() model = TestModel() trainer = Trainer( default_root_dir=tmpdir, fast_dev_run=True, callbacks=ck, gpus=1, strategy=DeepSpeedStrategy(logging_batch_size_per_gpu=value, zero_optimization=False), ) with pytest.raises(SystemExit): trainer.fit(model)
def test_deepspeed_with_invalid_config_path(tmpdir): """Test to ensure if we pass an invalid config path we throw an exception.""" with pytest.raises( MisconfigurationException, match="You passed in a path to a DeepSpeed config but the path does not exist" ): DeepSpeedStrategy(config="invalid_path.json")
def test_deepspeed_multigpu_partial_partition_parameters(tmpdir): """Test to ensure that a module that defines a layer inside the ``__init__`` and ``configure_sharded_model`` correctly converts all parameters to float16 when ``precision=16`` and runs successfully.""" class TestModel(ModelParallelBoringModel): def __init__(self): super().__init__() self.layer_2 = torch.nn.Linear(32, 32) def configure_sharded_model(self) -> None: self.layer = torch.nn.Linear(32, 2) def forward(self, x): x = self.layer_2(x) return self.layer(x) def on_train_epoch_start(self) -> None: assert all([x.dtype == torch.float16 for x in self.parameters()]) model = TestModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=1, fast_dev_run=True, precision=16, ) trainer.fit(model)
def test_deepspeed_summary(tmpdir): """Test to ensure that the summary contains the correct values when stage 3 is enabled and that the trainer enables the `DeepSpeedSummary` when DeepSpeed is used.""" model = BoringModel() total_parameters = sum(x.numel() for x in model.parameters()) class TestCallback(Callback): def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None: model_summary = DeepSpeedSummary(pl_module, max_depth=1) assert model_summary.total_parameters == total_parameters assert model_summary.trainable_parameters == total_parameters # check the additional params per device summary_data = model_summary._get_summary_data() params_per_device = summary_data[-1][-1] assert int( params_per_device[0]) == (model_summary.total_parameters // 2) trainer = Trainer( strategy=DeepSpeedStrategy(stage=3), default_root_dir=tmpdir, accelerator="gpu", fast_dev_run=True, devices=2, precision=16, enable_model_summary=True, callbacks=[TestCallback()], ) trainer.fit(model)
def test_deepspeed_config(tmpdir, deepspeed_zero_config): """Test to ensure deepspeed works correctly when passed a DeepSpeed config object including optimizers/schedulers and saves the model weights to load correctly.""" class TestCB(Callback): def on_train_start(self, trainer, pl_module) -> None: from deepspeed.runtime.lr_schedules import WarmupLR from deepspeed.runtime.zero.stage2 import FP16_DeepSpeedZeroOptimizer assert isinstance(trainer.optimizers[0], FP16_DeepSpeedZeroOptimizer) assert isinstance(trainer.optimizers[0].optimizer, torch.optim.SGD) assert isinstance(trainer.lr_scheduler_configs[0].scheduler, WarmupLR) model = BoringModel() trainer = Trainer( strategy=DeepSpeedStrategy(config=deepspeed_zero_config), default_root_dir=tmpdir, gpus=1, fast_dev_run=True, precision=16, callbacks=[TestCB()], ) trainer.fit(model) trainer.test(model)
def test_deepspeed_custom_precision_params(tmpdir): """Ensure if we modify the FP16 parameters via the DeepSpeedStrategy, the deepspeed config contains these changes.""" class TestCB(Callback): def on_train_start(self, trainer, pl_module) -> None: assert trainer.strategy.config["fp16"]["loss_scale"] == 10 assert trainer.strategy.config["fp16"]["initial_scale_power"] == 10 assert trainer.strategy.config["fp16"]["loss_scale_window"] == 10 assert trainer.strategy.config["fp16"]["hysteresis"] == 10 assert trainer.strategy.config["fp16"]["min_loss_scale"] == 10 raise SystemExit() model = BoringModel() ds = DeepSpeedStrategy(loss_scale=10, initial_scale_power=10, loss_scale_window=10, hysteresis=10, min_loss_scale=10) trainer = Trainer(default_root_dir=tmpdir, strategy=ds, precision=16, accelerator="gpu", devices=1, callbacks=[TestCB()]) with pytest.raises(SystemExit): trainer.fit(model)
def test_deepspeed_config(tmpdir, deepspeed_zero_config): """Test to ensure deepspeed works correctly when passed a DeepSpeed config object including optimizers/schedulers and saves the model weights to load correctly.""" class TestCB(Callback): def on_train_start(self, trainer, pl_module) -> None: from deepspeed.runtime.lr_schedules import WarmupLR from deepspeed.runtime.zero.stage2 import FP16_DeepSpeedZeroOptimizer assert isinstance(trainer.optimizers[0], FP16_DeepSpeedZeroOptimizer) assert isinstance(trainer.optimizers[0].optimizer, torch.optim.SGD) assert isinstance(trainer.lr_scheduler_configs[0].scheduler, WarmupLR) assert trainer.lr_scheduler_configs[0].interval == "step" assert trainer.lr_scheduler_configs[0].opt_idx == 0 model = BoringModel() lr_monitor = LearningRateMonitor() trainer = Trainer( strategy=DeepSpeedStrategy(config=deepspeed_zero_config), default_root_dir=tmpdir, gpus=1, log_every_n_steps=1, limit_train_batches=4, limit_val_batches=4, limit_test_batches=4, max_epochs=2, precision=16, callbacks=[TestCB(), lr_monitor], ) trainer.fit(model) trainer.test(model) assert list(lr_monitor.lrs) == ["lr-SGD"] assert len(set(lr_monitor.lrs["lr-SGD"])) == 8
def test_deepspeed_setup_train_dataloader(tmpdir): """Test DeepSpeed works when setup is required to call in the DataModule.""" class TestSetupIsCalledDataModule(LightningDataModule): def __init__(self): super().__init__() self._setup = False def setup(self, stage: Optional[str] = None) -> None: self._setup = True def train_dataloader(self): assert self._setup return DataLoader(RandomDataset(32, 64), batch_size=2) def val_dataloader(self): assert self._setup return DataLoader(RandomDataset(32, 64), batch_size=2) def test_dataloader(self): assert self._setup return DataLoader(RandomDataset(32, 64), batch_size=2) model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(logging_level=logging.INFO), gpus=1, fast_dev_run=True, ) dm = TestSetupIsCalledDataModule() with mock.patch("deepspeed.utils.logging.logger.warning", autospec=True) as mock_object: trainer.fit(model, datamodule=dm) assert any("Tried to infer the batch size" in str(arg) for arg in mock_object.call_args_list)
def test_deepspeed_multigpu_test(tmpdir): """Test to ensure we can use DeepSpeed with just test using ZeRO Stage 3.""" model = ModelParallelBoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) trainer.test(model)
def test_deepspeed_multigpu_stage_2_accumulated_grad_batches(tmpdir, offload_optimizer): """Test to ensure with Stage 2 and multiple GPUs, accumulated grad batches works.""" seed_everything(42) class VerificationCallback(Callback): def __init__(self): self.on_train_batch_start_called = False def on_train_batch_start(self, trainer, pl_module: LightningModule, batch: Any, batch_idx: int) -> None: deepspeed_engine = trainer.strategy.model assert trainer.global_step == deepspeed_engine.global_steps self.on_train_batch_start_called = True model = ModelParallelClassificationModel() dm = ClassifDataModule() verification_callback = VerificationCallback() trainer = Trainer( default_root_dir=tmpdir, enable_progress_bar=False, # TODO: this test fails with max_epochs >1 as there are leftover batches per epoch. # there's divergence in how Lightning handles the last batch of the epoch with how DeepSpeed does it. # we step the optimizers on the last batch but DeepSpeed keeps the accumulation for the next epoch max_epochs=1, strategy=DeepSpeedStrategy(stage=2, offload_optimizer=offload_optimizer), gpus=2, limit_train_batches=5, limit_val_batches=2, precision=16, accumulate_grad_batches=2, callbacks=[verification_callback], ) assert trainer.limit_train_batches % trainer.accumulate_grad_batches != 0, "leftover batches should be tested" trainer.fit(model, datamodule=dm) assert verification_callback.on_train_batch_start_called
def test_deepspeed_custom_activation_checkpointing_params_forwarded(tmpdir): """Ensure if we modify the activation checkpointing parameters, we pass these to deepspeed.checkpointing.configure correctly.""" ds = DeepSpeedStrategy( partition_activations=True, cpu_checkpointing=True, contiguous_memory_optimization=True, synchronize_checkpoint_boundary=True, ) model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, enable_progress_bar=False, fast_dev_run=1, strategy=ds, precision=16, accelerator="gpu", devices=1, ) with mock.patch("deepspeed.checkpointing.configure", wraps=deepspeed.checkpointing.configure ) as deepspeed_checkpointing_configure: trainer.fit(model) deepspeed_checkpointing_configure.assert_called_with( mpu_=None, partition_activations=True, contiguous_checkpointing=True, checkpoint_in_cpu=True, profile=None)
def test_deepspeed_multigpu_stage_3_warns_resume_training(tmpdir): """Test to ensure with Stage 3 and multiple GPUs that we can resume from training, throwing a warning that the optimizer state and scheduler states cannot be restored.""" dm = ClassifDataModule() model = BoringModel() checkpoint_path = os.path.join(tmpdir, "model.pt") trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True) trainer.fit(model) trainer.save_checkpoint(checkpoint_path) trainer = Trainer( default_root_dir=tmpdir, fast_dev_run=True, strategy=DeepSpeedStrategy(stage=3, load_full_weights=True), accelerator="gpu", devices=1, precision=16, ) with pytest.warns( UserWarning, match= "A single checkpoint file has been given. This means optimizer states cannot be restored. " "If you'd like to restore these states, you must " "provide a path to the originally saved DeepSpeed checkpoint.", ): trainer.fit(model, datamodule=dm, ckpt_path=checkpoint_path)
def test_deepspeed_strategy_env_variables(mock_deepspeed_distributed, tmpdir, platform): """Test to ensure that we setup distributed communication using correctly. When using windows, ranks environment variables should not be set, and deepspeed should handle this. """ trainer = Trainer(default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3)) strategy = trainer.strategy assert isinstance(strategy, DeepSpeedStrategy) with mock.patch("platform.system", return_value=platform) as mock_platform: strategy._init_deepspeed_distributed() mock_deepspeed_distributed.assert_called() mock_platform.assert_called() if platform == "Windows": # assert no env variables have been set within the DeepSpeedStrategy assert all(k not in os.environ for k in ("MASTER_PORT", "MASTER_ADDR", "RANK", "WORLD_SIZE", "LOCAL_RANK")) else: assert os.environ["MASTER_ADDR"] == str( trainer.strategy.cluster_environment.main_address) assert os.environ["MASTER_PORT"] == str( trainer.strategy.cluster_environment.main_port) assert os.environ["RANK"] == str(trainer.strategy.global_rank) assert os.environ["WORLD_SIZE"] == str(trainer.strategy.world_size) assert os.environ["LOCAL_RANK"] == str(trainer.strategy.local_rank)
def select_strategy(self) -> Strategy: if isinstance(self.distributed_backend, Accelerator) and self.distributed_backend.strategy is not None: plugin = self.distributed_backend.strategy elif self.use_ddp2: plugin = DDP2Strategy(parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment) elif self.use_ddp and self.use_deepspeed: plugin = DeepSpeedStrategy( cluster_environment=self.select_cluster_environment(), parallel_devices=self.parallel_devices ) elif self.use_ddp: use_slurm_ddp = self.use_ddp and self._is_slurm_managing_tasks() use_torchelastic_ddp = self.use_ddp and TorchElasticEnvironment.detect() use_kubeflow_ddp = self.use_ddp and KubeflowEnvironment.detect() use_ddp_spawn = self._strategy_type == _StrategyType.DDP_SPAWN use_ddp_cpu_spawn = use_ddp_spawn and self.use_cpu use_tpu_spawn = self.use_tpu and self._strategy_type == _StrategyType.TPU_SPAWN use_ddp_cpu_torch_elastic = use_ddp_cpu_spawn and TorchElasticEnvironment.detect() use_ddp_cpu_kubeflow = use_ddp_cpu_spawn and KubeflowEnvironment.detect() use_ddp_cpu_slurm = use_ddp_cpu_spawn and self._is_slurm_managing_tasks() use_ddp_sharded = self._strategy_type == _StrategyType.DDP_SHARDED use_ddp_sharded_spawn = self._strategy_type == _StrategyType.DDP_SHARDED_SPAWN use_ddp_fully_sharded = self._strategy_type == _StrategyType.DDP_FULLY_SHARDED if use_tpu_spawn: ddp_strategy_cls = TPUSpawnStrategy elif use_ddp_sharded: ddp_strategy_cls = DDPShardedStrategy elif use_ddp_sharded_spawn: ddp_strategy_cls = DDPSpawnShardedStrategy elif ( use_ddp_cpu_slurm or use_slurm_ddp or use_ddp_cpu_torch_elastic or use_torchelastic_ddp or use_kubeflow_ddp or use_ddp_cpu_kubeflow ): ddp_strategy_cls = DDPStrategy elif use_ddp_spawn or use_ddp_cpu_spawn: ddp_strategy_cls = DDPSpawnStrategy elif use_ddp_fully_sharded: ddp_strategy_cls = DDPFullyShardedStrategy else: ddp_strategy_cls = DDPStrategy plugin = ddp_strategy_cls( parallel_devices=self.parallel_devices, cluster_environment=self.cluster_environment ) elif self.use_dp: plugin = DataParallelStrategy(parallel_devices=self.parallel_devices) elif self.use_horovod: plugin = HorovodStrategy(parallel_devices=self.parallel_devices) elif self.use_tpu and isinstance(self.tpu_cores, list): plugin = SingleTPUStrategy(self.tpu_id) elif self.use_ipu: plugin = IPUStrategy(parallel_devices=self.parallel_devices) else: single_gpu_ordinal = device_parser.determine_root_gpu_device(self.parallel_device_ids) plugin = SingleDeviceStrategy(device=single_gpu_ordinal if self.use_gpu else "cpu") return plugin
def test_deepspeed_with_env_path(tmpdir, monkeypatch, deepspeed_config): """Test to ensure if we pass an env variable, we load the config from the path.""" config_path = os.path.join(tmpdir, "temp.json") with open(config_path, "w") as f: f.write(json.dumps(deepspeed_config)) monkeypatch.setenv("PL_DEEPSPEED_CONFIG_PATH", config_path) strategy = DeepSpeedStrategy() assert strategy.config == deepspeed_config
def test_deepspeed_multigpu_no_schedulers(tmpdir): """Test to ensure ZeRO Stage 3 works with a parallel model and no schedulers.""" model = ModelParallelBoringModelNoSchedulers() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) trainer.fit(model) _assert_save_model_is_equal(model, tmpdir, trainer)
def test_deepspeed_skip_backward_raises(tmpdir): class TestModel(BoringModel): def training_step(self, batch, batch_idx): return None model = TestModel() trainer = Trainer(default_root_dir=tmpdir, strategy=DeepSpeedStrategy(), gpus=1, fast_dev_run=True, precision=16) with pytest.raises(MisconfigurationException, match="returning `None` .* is not supported"): trainer.fit(model)
def test_deepspeed_with_meta_device(tmpdir): with init_meta_context(): model = BoringModel() assert model.layer.weight.device.type == "meta" trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) trainer.fit(model) assert model.layer.weight.device.type == "cpu"
def test_deepspeed_multigpu_stage_3(tmpdir, deepspeed_config): """Test to ensure ZeRO Stage 3 works with a parallel model.""" model = ModelParallelBoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) trainer.fit(model) trainer.test(model) _assert_save_model_is_equal(model, tmpdir, trainer)
def test_deepspeed_multigpu(tmpdir): """Test to ensure that DeepSpeed with multiple GPUs works and deepspeed distributed is initialized correctly.""" model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) with mock.patch("deepspeed.init_distributed", wraps=deepspeed.init_distributed) as mock_deepspeed_distributed: trainer.fit(model) mock_deepspeed_distributed.assert_called_once() trainer.test(model) _assert_save_model_is_equal(model, tmpdir, trainer)
def test_deepspeed_custom_activation_checkpointing_params(tmpdir): """Ensure if we modify the activation checkpointing parameters, the deepspeed config contains these changes.""" ds = DeepSpeedStrategy( partition_activations=True, cpu_checkpointing=True, contiguous_memory_optimization=True, synchronize_checkpoint_boundary=True, ) checkpoint_config = ds.config["activation_checkpointing"] assert checkpoint_config["partition_activations"] assert checkpoint_config["cpu_checkpointing"] assert checkpoint_config["contiguous_memory_optimization"] assert checkpoint_config["synchronize_checkpoint_boundary"]
def test_deepspeed_stage_3_save_warning(tmpdir): """Test to ensure that DeepSpeed Stage 3 gives a warning when saving on rank zero.""" model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=2, fast_dev_run=True, precision=16 ) trainer.fit(model) checkpoint_path = os.path.join(tmpdir, "model.pt") # both ranks need to call save checkpoint, however only rank 0 needs to check the warning context_manager = ( pytest.warns(UserWarning, match="each worker will save a shard of the checkpoint within a directory.") if trainer.is_global_zero else contextlib.suppress() ) with context_manager: trainer.save_checkpoint(checkpoint_path)
def test_deepspeed_multigpu_stage_3_manual_optimization( tmpdir, deepspeed_config): """Test to ensure ZeRO Stage 3 works with a parallel model.""" model = ModelParallelBoringModelManualOptim() model.training_epoch_end = None trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=2, fast_dev_run=True, precision=16, ) trainer.fit(model) trainer.test(model) _assert_save_model_is_equal(model, tmpdir, trainer)
def test_warn_deepspeed_ignored(tmpdir): class TestModel(BoringModel): def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None: return loss.backward() model = TestModel() trainer = Trainer( fast_dev_run=True, default_root_dir=tmpdir, strategy=DeepSpeedStrategy(), gpus=1, precision=16, track_grad_norm=2, ) from pytorch_lightning.plugins.precision.deepspeed import warning_cache with pytest.warns(UserWarning, match="will be ignored since DeepSpeed handles the backward"): trainer.fit(model) assert any("track_grad_norm=2.0)' but this is not supported" in w for w in warning_cache)
def test_deepspeed_multigpu_test_rnn(tmpdir): """Test to ensure that turning off explicit partitioning of the entire module for ZeRO Stage 3 works when training with certain layers which will crash with explicit partitioning.""" class TestModel(BoringModel): def __init__(self): super().__init__() self.rnn = torch.nn.GRU(32, 32) def on_train_epoch_start(self) -> None: assert all([x.dtype == torch.float16 for x in self.parameters()]) model = TestModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), gpus=1, fast_dev_run=True, precision=16, ) trainer.fit(model)
def test_deepspeed_collate_checkpoint(tmpdir): """Test to ensure that with DeepSpeed Stage 3 we can collate the sharded checkpoints into a single file.""" model = BoringModel() trainer = Trainer( default_root_dir=tmpdir, strategy=DeepSpeedStrategy(stage=3), accelerator="gpu", devices=2, fast_dev_run=True, precision=16, ) trainer.fit(model) checkpoint_path = os.path.join(tmpdir, "model.pt") checkpoint_path = trainer.strategy.broadcast(checkpoint_path) trainer.save_checkpoint(checkpoint_path) trainer.strategy.barrier() if trainer.is_global_zero: # ensure function call works output_path = os.path.join(tmpdir, "single_model.pt") convert_zero_checkpoint_to_fp32_state_dict(checkpoint_path, output_path) _assert_checkpoint_equal(model, output_path)