Check if the CustomAllreduce function works properly.
(self)
| 43 | fleet.init(is_collective=True, strategy=strategy) |
| 44 | |
| 45 | def test_case(self): |
| 46 | """ |
| 47 | Check if the CustomAllreduce function works properly. |
| 48 | """ |
| 49 | |
| 50 | mns = [[1, 2048], [2, 4096], [20, 4096], [128, 4096], [256, 4096], [256, 8192]] |
| 51 | |
| 52 | hcg = fleet.get_hybrid_communicate_group() |
| 53 | model_parallel_group = hcg.get_model_parallel_group() |
| 54 | fa = CustomAllreduce(model_parallel_group) |
| 55 | |
| 56 | for m, n in mns: |
| 57 | data_custom_ar = paddle.rand([m, n], dtype="bfloat16") |
| 58 | data_paddle = data_custom_ar.clone() |
| 59 | if fa.should_custom_ar(data_custom_ar): |
| 60 | data_custom_ar = fa.custom_all_reduce(data_custom_ar) |
| 61 | dist.all_reduce(data_paddle) |
| 62 | if dist.get_rank() == 0: |
| 63 | np.testing.assert_allclose( |
| 64 | data_custom_ar.numpy(), |
| 65 | data_paddle.numpy(), |
| 66 | rtol=1e-04, |
| 67 | atol=1e-04, |
| 68 | ) |
| 69 | |
| 70 | |
| 71 | if __name__ == "__main__": |
nothing calls this directly
no test coverage detected