已合并
PREMUL_SUM dtype constraints #34518
jizewei创建于 4月27日
PREMUL_SUM dtype constraints #34518
已合并
共 5 个文件变更+14-13
| @@ -170,12 +170,10 @@ class HcomAllReduceTest(TestCase): | |||
| 170 | 170 | ||
| 171 | def test_dist_all_reduce_pre_mul(self): | 171 | def test_dist_all_reduce_pre_mul(self): |
| 172 | ranks = [2] | 172 | ranks = [2] |
| 173 | - dtype_list = [np.int32, np.int8] | 173 | + dtype_list = [np.float32, np.float16] |
| 174 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] | 174 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] |
| 175 | for world_size in ranks: | 175 | for world_size in ranks: |
| 176 | for shape in shape_format: | 176 | for shape in shape_format: |
| 177 | - if shape[0] == np.int8: | ||
| 178 | - shape[1] = 0 | ||
| 179 | exp_input, input1 = create_common_tensor(shape, -10, 10) | 177 | exp_input, input1 = create_common_tensor(shape, -10, 10) |
| 180 | expected = self._construct_excepted_result(exp_input, world_size, shape[0], dist.ReduceOp.SUM) | 178 | expected = self._construct_excepted_result(exp_input, world_size, shape[0], dist.ReduceOp.SUM) |
| 181 | expected = expected * 2 | 179 | expected = expected * 2 |
| @@ -128,12 +128,10 @@ class HcclReduceTest(TestCase): | |||
| 128 | 128 | ||
| 129 | def test_reduce_dist_pre_mul(self): | 129 | def test_reduce_dist_pre_mul(self): |
| 130 | ranks = [2] | 130 | ranks = [2] |
| 131 | - dtype_list = [np.int32, np.int8] | 131 | + dtype_list = [np.float32, np.float16] |
| 132 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] | 132 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] |
| 133 | for world_size in ranks: | 133 | for world_size in ranks: |
| 134 | for shape in shape_format: | 134 | for shape in shape_format: |
| 135 | - if shape[0] == np.int8: | ||
| 136 | - shape[1] = 0 | ||
| 137 | exp_input, input1 = create_common_tensor(shape, -10, 10) | 135 | exp_input, input1 = create_common_tensor(shape, -10, 10) |
| 138 | expected = self._construct_excepted_result(exp_input, world_size, shape[0], dist.ReduceOp.SUM) | 136 | expected = self._construct_excepted_result(exp_input, world_size, shape[0], dist.ReduceOp.SUM) |
| 139 | expected = expected * 2 | 137 | expected = expected * 2 |
| @@ -181,13 +181,11 @@ class HcclReduceScatterTest(HcclReduceScatterTestBase): | |||
| 181 | 181 | ||
| 182 | def test_reduce_scatter_pre_mul(self): | 182 | def test_reduce_scatter_pre_mul(self): |
| 183 | ranks = [2] | 183 | ranks = [2] |
| 184 | - dtype_list = [np.int32, np.int8] | 184 | + dtype_list = [np.float32, np.float16] |
| 185 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] | 185 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] |
| 186 | 186 | ||
| 187 | for world_size in ranks: | 187 | for world_size in ranks: |
| 188 | for shape in shape_format: | 188 | for shape in shape_format: |
| 189 | - if shape[0] == np.int8: | ||
| 190 | - shape[1] = 0 | ||
| 191 | input_list = [] | 189 | input_list = [] |
| 192 | for _ in range(world_size): | 190 | for _ in range(world_size): |
| 193 | _, input1 = create_common_tensor(shape, -10, 10) | 191 | _, input1 = create_common_tensor(shape, -10, 10) |
| @@ -86,12 +86,10 @@ class HcclReduceScatterBaseTest(HcclReduceScatterTestBase): | |||
| 86 | 86 | ||
| 87 | def test_reduce_scatter_base_pre_mul(self): | 87 | def test_reduce_scatter_base_pre_mul(self): |
| 88 | ranks = [2] | 88 | ranks = [2] |
| 89 | - dtype_list = [np.int32, np.int8] | 89 | + dtype_list = [np.float32, np.float16] |
| 90 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] | 90 | shape_format = [[i, 2, [4, 9]] for i in dtype_list] |
| 91 | for world_size in ranks: | 91 | for world_size in ranks: |
| 92 | for shape in shape_format: | 92 | for shape in shape_format: |
| 93 | - if shape[0] == np.int8: | ||
| 94 | - shape[1] = 0 | ||
| 95 | input_list = [] | 93 | input_list = [] |
| 96 | for _ in range(world_size): | 94 | for _ in range(world_size): |
| 97 | _, input1 = create_common_tensor(shape, -10, 10) | 95 | _, input1 = create_common_tensor(shape, -10, 10) |
| @@ -133,12 +133,21 @@ uint64_t getNumelForHCCL(const at::Tensor& self) | |||
| 133 | 133 | ||
| 134 | HcclReduceOp getHcclReduceOp(const c10d::ReduceOp reduceOp, at::Tensor& input) | 134 | HcclReduceOp getHcclReduceOp(const c10d::ReduceOp reduceOp, at::Tensor& input) |
| 135 | { | 135 | { |
| 136 | - if (reduceOp == c10d::ReduceOp::AVG || reduceOp == c10d::ReduceOp::PREMUL_SUM) { | 136 | + if (reduceOp == c10d::ReduceOp::AVG) { |
| 137 | // HCCL does not support ReduceOp::AVG yet | 137 | // HCCL does not support ReduceOp::AVG yet |
| 138 | // PTA supports it by summing first, then dividing | 138 | // PTA supports it by summing first, then dividing |
| 139 | return HCCL_REDUCE_SUM; | 139 | return HCCL_REDUCE_SUM; |
| 140 | } | 140 | } |
| 141 | 141 | ||
| 142 | + if (reduceOp == c10d::ReduceOp::PREMUL_SUM) { | ||
| 143 | + TORCH_CHECK( | ||
| 144 | + input.scalar_type() == at::kHalf || input.scalar_type() == at::kFloat || | ||
| 145 | + input.scalar_type() == at::kBFloat16 || input.scalar_type() == at::kDouble, | ||
| 146 | + "PreMulSum Data type must be half, float, bfloat16 or double", | ||
| 147 | + DIST_ERROR(ErrCode::TYPE)); | ||
| 148 | + return HCCL_REDUCE_SUM; | ||
| 149 | + } | ||
| 150 | + | ||
| 142 | if (reduceOp == c10d::ReduceOp::SUM && input.scalar_type() == at::kBool) { | 151 | if (reduceOp == c10d::ReduceOp::SUM && input.scalar_type() == at::kBool) { |
| 143 | // For bool tensors, map sum to max, which both represent a bitwise or. | 152 | // For bool tensors, map sum to max, which both represent a bitwise or. |
| 144 | // This is to prevent overflow issues with sum, since we use uint8 to | 153 | // This is to prevent overflow issues with sum, since we use uint8 to |