| 239 | self.has_moe_params = True |
| 240 | |
| 241 | def check_using_norm(self, norm_group, reduce_overflow=True): |
| 242 | # TODO: I don't think reduce_overflow is needed if mpu is None |
| 243 | overflow = -1 in norm_group |
| 244 | overflow_gpu = get_accelerator().FloatTensor([overflow]) |
| 245 | if self.has_moe_params: |
| 246 | # In this case, we need to do an all_reduce across |
| 247 | # the expert_parallel_group, so that if there was |
| 248 | # an overflow due to expert weights, we detect it |
| 249 | |
| 250 | # Only need to check groups.get_largest_expert_parallel_group() |
| 251 | dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=groups._get_max_expert_parallel_group()) |
| 252 | if self.mpu is not None: |
| 253 | dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX, group=self.mpu.get_model_parallel_group()) |
| 254 | elif reduce_overflow: |
| 255 | dist.all_reduce(overflow_gpu, op=dist.ReduceOp.MAX) |
| 256 | dist.barrier() |
| 257 | overflow = overflow_gpu[0].item() |
| 258 | return bool(overflow) |
| 259 | |
| 260 | def check(self, param_groups=None): |
| 261 | params = [] |