diff --git a/tests/tests_fabric/strategies/test_deepspeed_integration.py b/tests/tests_fabric/strategies/test_deepspeed_integration.py index 6cca479b65..9f4923e987 100644 --- a/tests/tests_fabric/strategies/test_deepspeed_integration.py +++ b/tests/tests_fabric/strategies/test_deepspeed_integration.py @@ -31,78 +31,79 @@ from tests_fabric.test_fabric import BoringModel @RunIf(min_cuda_gpus=2, standalone=True, deepspeed=True) def test_deepspeed_multiple_models(): - def run(fabric_obj): - model = BoringModel() - optimizer = torch.optim.SGD(model.parameters(), lr=0.0001) - model, optimizer = fabric_obj.setup(model, optimizer) - - for i in range(2): - optimizer.zero_grad() - x = model(torch.randn(1, 32).to(fabric_obj.device)) - loss = x.sum() - if i == 0: - # the weights are not initialized with stage 3 until backward is run once - assert all(w.nelement() == 0 for w in model.state_dict().values()) - fabric_obj.backward(loss, model=model) - if i == 0: - # save for later to check that the weights were updated - state_dict = deepcopy(model.state_dict()) - optimizer.step() - - # check that the model trained, the weights from step 1 do not match the weights from step 2 - for mw_b, mw_a in zip(state_dict.values(), model.state_dict().values()): - assert not torch.allclose(mw_b, mw_a) - - fabric_obj.seed_everything(42) - model_1 = BoringModel() - optimizer_1 = torch.optim.SGD(model_1.parameters(), lr=0.0001) - - fabric_obj.seed_everything(42) - model_2 = BoringModel() - optimizer_2 = torch.optim.SGD(model_2.parameters(), lr=0.0001) - - for mw_1, mw_2 in zip(model_1.state_dict().values(), model_2.state_dict().values()): - assert torch.allclose(mw_1, mw_2) - - model_1, optimizer_1 = fabric_obj.setup(model_1, optimizer_1) - model_2, optimizer_2 = fabric_obj.setup(model_2, optimizer_2) - - # train model_1 first - fabric_obj.seed_everything(42) - data_list = [] - for _ in range(2): - optimizer_1.zero_grad() - data = torch.randn(1, 32).to(fabric_obj.device) - data_list.append(data) - x = model_1(data) - loss = x.sum() - fabric_obj.backward(loss, model=model_1) - optimizer_1.step() - - # the weights do not match - assert all(w.nelement() > 1 for w in model_1.state_dict().values()) - assert all(w.nelement() == 0 for w in model_2.state_dict().values()) - - # now train model_2 with the same data - for data in data_list: - optimizer_2.zero_grad() - x = model_2(data) - loss = x.sum() - fabric_obj.backward(loss, model=model_2) - optimizer_2.step() - - # the weights should match - for mw_1, mw_2 in zip(model_1.state_dict().values(), model_2.state_dict().values()): - assert torch.allclose(mw_1, mw_2) - - # Verify collectives works as expected - ranks = fabric_obj.all_gather(torch.tensor([fabric_obj.local_rank]).to(fabric_obj.device)) - assert torch.allclose(ranks.cpu(), torch.tensor([[0], [1]])) - assert fabric_obj.broadcast(True) - assert fabric_obj.is_global_zero == (fabric_obj.local_rank == 0) - fabric = Fabric(strategy=DeepSpeedStrategy(stage=3, logging_batch_size_per_gpu=1), devices=2, accelerator="gpu") - fabric.launch(run) + fabric.launch() + + with fabric.init_module(): + model = BoringModel() + + optimizer = torch.optim.SGD(model.parameters(), lr=0.0001) + model, optimizer = fabric.setup(model, optimizer) + + for i in range(2): + optimizer.zero_grad() + x = model(torch.randn(1, 32).to(fabric.device)) + loss = x.sum() + if i == 0: + # the weights are not initialized with stage 3 until backward is run once + assert all(w.nelement() == 0 for w in model.state_dict().values()) + fabric.backward(loss, model=model) + if i == 0: + # save for later to check that the weights were updated + state_dict = deepcopy(model.state_dict()) + optimizer.step() + + # check that the model trained, the weights from step 1 do not match the weights from step 2 + for mw_b, mw_a in zip(state_dict.values(), model.state_dict().values()): + assert not torch.allclose(mw_b, mw_a) + + fabric.seed_everything(42) + model_1 = BoringModel() + optimizer_1 = torch.optim.SGD(model_1.parameters(), lr=0.0001) + + fabric.seed_everything(42) + model_2 = BoringModel() + optimizer_2 = torch.optim.SGD(model_2.parameters(), lr=0.0001) + + for mw_1, mw_2 in zip(model_1.state_dict().values(), model_2.state_dict().values()): + assert torch.allclose(mw_1, mw_2) + + model_1, optimizer_1 = fabric.setup(model_1, optimizer_1) + model_2, optimizer_2 = fabric.setup(model_2, optimizer_2) + + # train model_1 first + fabric.seed_everything(42) + data_list = [] + for _ in range(2): + optimizer_1.zero_grad() + data = torch.randn(1, 32).to(fabric.device) + data_list.append(data) + x = model_1(data) + loss = x.sum() + fabric.backward(loss, model=model_1) + optimizer_1.step() + + # the weights do not match + assert all(w.nelement() > 1 for w in model_1.state_dict().values()) + assert all(w.nelement() == 0 for w in model_2.state_dict().values()) + + # now train model_2 with the same data + for data in data_list: + optimizer_2.zero_grad() + x = model_2(data) + loss = x.sum() + fabric.backward(loss, model=model_2) + optimizer_2.step() + + # the weights should match + for mw_1, mw_2 in zip(model_1.state_dict().values(), model_2.state_dict().values()): + assert torch.allclose(mw_1, mw_2) + + # Verify collectives works as expected + ranks = fabric.all_gather(torch.tensor([fabric.local_rank]).to(fabric.device)) + assert torch.allclose(ranks.cpu(), torch.tensor([[0], [1]])) + assert fabric.broadcast(True) + assert fabric.is_global_zero == (fabric.local_rank == 0) @RunIf(min_cuda_gpus=1, deepspeed=True) @@ -117,19 +118,16 @@ def test_deepspeed_multiple_models(): ) def test_deepspeed_auto_batch_size_config_select(dataset_cls, logging_batch_size_per_gpu, expected_batch_size): """Test to ensure that the batch size is correctly set as expected for deepspeed logging purposes.""" - - def run(fabric_obj): - assert isinstance(fabric_obj._strategy, DeepSpeedStrategy) - _ = fabric_obj.setup_dataloaders(DataLoader(dataset_cls(32, 64))) - config = fabric_obj._strategy.config - assert config["train_micro_batch_size_per_gpu"] == expected_batch_size - fabric = Fabric( accelerator="cuda", devices=1, strategy=DeepSpeedStrategy(logging_batch_size_per_gpu=logging_batch_size_per_gpu, zero_optimization=False), ) - fabric.launch(run) + fabric.launch() + assert isinstance(fabric._strategy, DeepSpeedStrategy) + _ = fabric.setup_dataloaders(DataLoader(dataset_cls(32, 64))) + config = fabric._strategy.config + assert config["train_micro_batch_size_per_gpu"] == expected_batch_size @RunIf(min_cuda_gpus=1, standalone=True, deepspeed=True) @@ -137,35 +135,24 @@ def test_deepspeed_configure_optimizers(): """Test that the deepspeed strategy with default initialization wraps the optimizer correctly.""" from deepspeed.runtime.zero.stage_1_and_2 import DeepSpeedZeroOptimizer - def run(fabric_obj): - model = nn.Linear(3, 3) - optimizer = torch.optim.SGD(model.parameters(), lr=0.01) - model, optimizer = fabric_obj.setup(model, optimizer) - assert isinstance(optimizer.optimizer, DeepSpeedZeroOptimizer) - assert isinstance(optimizer.optimizer.optimizer, torch.optim.SGD) - fabric = Fabric( strategy=DeepSpeedStrategy(), accelerator="cuda", devices=1, precision="16-mixed", ) - fabric.launch(run) + fabric.launch() + model = nn.Linear(3, 3) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + model, optimizer = fabric.setup(model, optimizer) + assert isinstance(optimizer.optimizer, DeepSpeedZeroOptimizer) + assert isinstance(optimizer.optimizer.optimizer, torch.optim.SGD) @RunIf(min_cuda_gpus=1, deepspeed=True) def test_deepspeed_custom_precision_params(): """Test that if the FP16 parameters are set via the DeepSpeedStrategy, the deepspeed config contains these changes.""" - - def run(fabric_obj): - assert fabric_obj._strategy._config_initialized - assert fabric_obj._strategy.config["fp16"]["loss_scale"] == 10 - assert fabric_obj._strategy.config["fp16"]["initial_scale_power"] == 11 - assert fabric_obj._strategy.config["fp16"]["loss_scale_window"] == 12 - assert fabric_obj._strategy.config["fp16"]["hysteresis"] == 13 - assert fabric_obj._strategy.config["fp16"]["min_loss_scale"] == 14 - strategy = DeepSpeedStrategy( loss_scale=10, initial_scale_power=11, loss_scale_window=12, hysteresis=13, min_loss_scale=14 ) @@ -175,7 +162,13 @@ def test_deepspeed_custom_precision_params(): accelerator="cuda", devices=1, ) - fabric.launch(run) + fabric.launch() + assert fabric._strategy._config_initialized + assert fabric._strategy.config["fp16"]["loss_scale"] == 10 + assert fabric._strategy.config["fp16"]["initial_scale_power"] == 11 + assert fabric._strategy.config["fp16"]["loss_scale_window"] == 12 + assert fabric._strategy.config["fp16"]["hysteresis"] == 13 + assert fabric._strategy.config["fp16"]["min_loss_scale"] == 14 @RunIf(min_cuda_gpus=1, standalone=True, deepspeed=True) @@ -184,21 +177,6 @@ def test_deepspeed_custom_activation_checkpointing_params_forwarded(): correctly.""" import deepspeed - def run(fabric_obj): - model = nn.Linear(3, 3) - optimizer = torch.optim.Adam(model.parameters()) - - with mock.patch("deepspeed.checkpointing.configure", wraps=deepspeed.checkpointing.configure) as configure: - fabric_obj.setup(model, optimizer) - - configure.assert_called_with( - mpu_=None, - partition_activations=True, - contiguous_checkpointing=True, - checkpoint_in_cpu=True, - profile=None, - ) - strategy = DeepSpeedStrategy( partition_activations=True, cpu_checkpointing=True, @@ -211,7 +189,20 @@ def test_deepspeed_custom_activation_checkpointing_params_forwarded(): accelerator="cuda", devices=1, ) - fabric.launch(run) + fabric.launch() + model = nn.Linear(3, 3) + optimizer = torch.optim.Adam(model.parameters()) + + with mock.patch("deepspeed.checkpointing.configure", wraps=deepspeed.checkpointing.configure) as configure: + fabric.setup(model, optimizer) + + configure.assert_called_with( + mpu_=None, + partition_activations=True, + contiguous_checkpointing=True, + checkpoint_in_cpu=True, + profile=None, + ) class ModelParallelClassification(BoringFabric): diff --git a/tests/tests_fabric/test_fabric.py b/tests/tests_fabric/test_fabric.py index 71ab649d5e..6d5fcef524 100644 --- a/tests/tests_fabric/test_fabric.py +++ b/tests/tests_fabric/test_fabric.py @@ -55,13 +55,17 @@ class BoringModel(nn.Module): def test_run_input_output(): """Test that the dynamically patched run() method receives the input arguments and returns the result.""" - def run(fabric_obj, *args, **kwargs): - fabric_obj.run_args = args - fabric_obj.run_kwargs = kwargs - return "result" + class RunFabric(Fabric): + run_args = () + run_kwargs = {} - fabric = Fabric() - result = fabric.launch(run, 1, 2, three=3) + def run(self, *args, **kwargs): + self.run_args = args + self.run_kwargs = kwargs + return "result" + + fabric = RunFabric() + result = fabric.run(1, 2, three=3) assert result == "result" assert fabric.run_args == (1, 2) assert fabric.run_kwargs == {"three": 3} @@ -315,8 +319,8 @@ def test_setup_dataloaders_return_type(): @mock.patch("lightning.fabric.fabric._replace_dunder_methods") def test_setup_dataloaders_captures_dataloader_arguments(ctx_manager): - """Test that Fabric intercepts the DataLoader constructor arguments with a context manager in its run - method.""" + """Test that Fabric intercepts the DataLoader constructor arguments with a context manager when launching a + function.""" def run(_): # One for BatchSampler, another for DataLoader @@ -534,26 +538,25 @@ def test_to_device(accelerator, expected): if not pjrt.using_pjrt(): expected = "xla:1" - def run(_): - expected_device = torch.device(expected) - - # module - module = torch.nn.Linear(2, 3) - module = fabric.to_device(module) - assert all(param.device == expected_device for param in module.parameters()) - - # tensor - tensor = torch.rand(2, 2) - tensor = fabric.to_device(tensor) - assert tensor.device == expected_device - - # collection - collection = {"data": torch.rand(2, 2), "int": 1} - collection = fabric.to_device(collection) - assert collection["data"].device == expected_device - fabric = Fabric(accelerator=accelerator, devices=1) - fabric.launch(run) + fabric.launch() + + expected_device = torch.device(expected) + + # module + module = torch.nn.Linear(2, 3) + module = fabric.to_device(module) + assert all(param.device == expected_device for param in module.parameters()) + + # tensor + tensor = torch.rand(2, 2) + tensor = fabric.to_device(tensor) + assert tensor.device == expected_device + + # collection + collection = {"data": torch.rand(2, 2), "int": 1} + collection = fabric.to_device(collection) + assert collection["data"].device == expected_device def test_rank_properties():