| 66 | |
| 67 | |
| 68 | class FP16_Module(MegatronModule): |
| 69 | def __init__(self, module): |
| 70 | super(FP16_Module, self).__init__() |
| 71 | self.add_module('module', module.half()) |
| 72 | |
| 73 | def forward(self, *inputs, **kwargs): |
| 74 | return fp16_to_fp32(self.module(*(fp32_to_fp16(inputs)), **kwargs)) |
| 75 | |
| 76 | def state_dict(self, destination=None, prefix='', keep_vars=False): |
| 77 | return self.module.state_dict(destination, prefix, keep_vars) |
| 78 | |
| 79 | def state_dict_for_save_checkpoint(self, destination=None, prefix='', |
| 80 | keep_vars=False): |
| 81 | return self.module.state_dict_for_save_checkpoint(destination, prefix, |
| 82 | keep_vars) |
| 83 | |
| 84 | def load_state_dict(self, state_dict, strict=True): |
| 85 | self.module.load_state_dict(state_dict, strict=strict) |
| 86 | |
| 87 | # TODO: Update overflow check + downscale to use Carl's fused kernel. |
| 88 | |