GPT OSS MXFP4 fix

This commit is contained in:
Daniel Han 2025-08-15 04:00:00 -07:00
commit d376e62038
2 changed files with 5 additions and 2 deletions

View file

@ -621,6 +621,7 @@ class FastModel(FastBaseModel):
"if hasattr(module, x): "\
"setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n"\
"x = 'down_proj_bias'\n"\
"setattr(module, x, torch.nn.Parameter(getattr(module, x).to(torch.float32)) if isinstance(getattr(module, x), torch.nn.Parameter) else getattr(module, x).to(torch.float32))\n"\
";"
else:
# Set down projection compute dtype to be float32 for float16 machines

View file

@ -457,8 +457,10 @@ class FastBaseModel:
# Edit data-types
if custom_datatype is not None:
for jj, (name, module) in enumerate(model.named_modules()):
exec(custom_datatype)
with torch.no_grad():
for jj, (name, module) in enumerate(model.named_modules()):
exec(custom_datatype)
pass
pass
pass
# Clear deleted GPU items