| 157 | |
| 158 | |
| 159 | class Float16Module(MegatronModule): |
| 160 | def __init__(self, module, args): |
| 161 | super(Float16Module, self).__init__() |
| 162 | |
| 163 | if args.fp16: |
| 164 | self.add_module("module", module.half()) |
| 165 | |
| 166 | def float16_convertor(val): |
| 167 | return val.half() |
| 168 | |
| 169 | elif args.bf16: |
| 170 | self.add_module("module", module.bfloat16()) |
| 171 | |
| 172 | def float16_convertor(val): |
| 173 | return val.bfloat16() |
| 174 | |
| 175 | else: |
| 176 | raise Exception("should not be here") |
| 177 | |
| 178 | self.float16_convertor = float16_convertor |
| 179 | |
| 180 | def forward(self, *inputs, **kwargs): |
| 181 | if mpu.is_pipeline_first_stage(): |
| 182 | inputs = fp32_to_float16(inputs, self.float16_convertor) |
| 183 | outputs = self.module(*inputs, **kwargs) |
| 184 | if mpu.is_pipeline_last_stage(): |
| 185 | outputs = float16_to_fp32(outputs) |
| 186 | return outputs |
| 187 | |
| 188 | def state_dict(self, destination=None, prefix="", keep_vars=False): |
| 189 | return self.module.state_dict(destination, prefix, keep_vars) |
| 190 | |
| 191 | def state_dict_for_save_checkpoint( |
| 192 | self, destination=None, prefix="", keep_vars=False |
| 193 | ): |
| 194 | return self.module.state_dict_for_save_checkpoint( |
| 195 | destination, prefix, keep_vars |
| 196 | ) |
| 197 | |
| 198 | def load_state_dict(self, state_dict, strict=True): |
| 199 | self.module.load_state_dict(state_dict, strict=strict) |