(self, sequential_model, simple_config, batch_input, activation_checkpoints, use_compile)
| 77 | @pytest.mark.parametrize("activation_checkpoints", [False, True]) |
| 78 | @pytest.mark.parametrize("use_compile", [False, True]) |
| 79 | def test(self, sequential_model, simple_config, batch_input, activation_checkpoints, use_compile): |
| 80 | base_model = copy.deepcopy(sequential_model) |
| 81 | base_input = batch_input.clone().detach() |
| 82 | base_output = base_model(base_input) |
| 83 | base_output = base_output |
| 84 | base_params = sum(p.numel() for p in base_model.parameters()) |
| 85 | |
| 86 | pipe_model = copy.deepcopy(sequential_model) |
| 87 | pipe_model = PipelineModule(layers=pipe_model, num_stages=2) |
| 88 | if (use_compile): |
| 89 | pipe_model.compile() |
| 90 | # Ensure all parameters are accounted for. |
| 91 | my_params = sum(p.numel() for p in pipe_model.parameters()) |
| 92 | total_pipe_params = torch.LongTensor([my_params]).to(get_accelerator().device_name()) |
| 93 | dist.all_reduce(total_pipe_params) |
| 94 | total_pipe_params = total_pipe_params.item() |
| 95 | assert total_pipe_params == base_params |
| 96 | |
| 97 | pipe_model, _, _, _ = deepspeed.initialize(config=simple_config, |
| 98 | model=pipe_model, |
| 99 | model_parameters=[p for p in pipe_model.parameters()]) |
| 100 | |
| 101 | if activation_checkpoints: |
| 102 | deepspeed.checkpointing.configure(None, |
| 103 | deepspeed_config=pipe_model.config, |
| 104 | partition_activations=True, |
| 105 | contiguous_checkpointing=True, |
| 106 | num_checkpoints=9) |
| 107 | |
| 108 | if pipe_model.is_first_stage or pipe_model.is_last_stage: |
| 109 | pipe_input = base_input.clone().detach().to(get_accelerator().device_name()) |
| 110 | # label 0 is meaningless |
| 111 | dataset = [(pipe_input, 0)] |
| 112 | loader = RepeatingLoader(dataset) |
| 113 | data_iter = iter(loader) |
| 114 | else: |
| 115 | data_iter = None |
| 116 | |
| 117 | pipe_output = pipe_model.eval_batch(data_iter=data_iter) |
| 118 | |
| 119 | base_output = base_output.to('cpu') |
| 120 | pipe_output = pipe_output.to('cpu') |
| 121 | |
| 122 | assert torch.allclose(base_output, pipe_output, atol=1e-4) |
| 123 | |
| 124 | |
| 125 | class TestPipeModuleCheckpointInterval(DistributedTest): |
nothing calls this directly
no test coverage detected