Pack MXFP4 nibbles along the row, not across groups

#2
by akopytko - opened
Files changed (1) hide show
  1. .meta/create_model.py +1 -0
.meta/create_model.py CHANGED
@@ -307,6 +307,7 @@ def create_model(source_model_id: str, source_revision: str, save_folder: str) -
307
  normalized = (blocks / scales.unsqueeze(-1)).clamp(-6, 6)
308
  codes = torch.bucketize(normalized.abs(), fp4_boundaries)
309
  codes += normalized.signbit() * 8
 
310
  state_dict[name.removesuffix(".weight") + ".weight_packed"] = (codes[:, ::2] | (codes[:, 1::2] << 4)).to(
311
  torch.uint8
312
  )
 
307
  normalized = (blocks / scales.unsqueeze(-1)).clamp(-6, 6)
308
  codes = torch.bucketize(normalized.abs(), fp4_boundaries)
309
  codes += normalized.signbit() * 8
310
+ codes = codes.view(out_features, in_features)
311
  state_dict[name.removesuffix(".weight") + ".weight_packed"] = (codes[:, ::2] | (codes[:, 1::2] << 4)).to(
312
  torch.uint8
313
  )