Merge pull request #22 from pytorch-labs/better_names

better buffer names

Author: vkuzo

float8_playground/float8_linear.py

@@ -33,15 +33,15 @@ def forward(
         x_fp8,
         w_fp8,
         b_fp8,
-        float8_amax_out,
-        float8_amax_dL_dX,
-        float8_amax_dL_dW,
-        float8_amax_dL_dY,
+        fp8_amax_y,
+        fp8_amax_dL_dX,
+        fp8_amax_dL_dW,
+        fp8_amax_dL_dY,
         fw_amax_initialized,
         bw_amax_initialized,
     ):
         ctx.save_for_backward(
-            x_fp8, w_fp8, b_fp8, float8_amax_dL_dX, float8_amax_dL_dW, float8_amax_dL_dY,
+            x_fp8, w_fp8, b_fp8, fp8_amax_dL_dX, fp8_amax_dL_dW, fp8_amax_dL_dY,
             bw_amax_initialized)
         orig_shape = x_fp8._data.shape
         x_fp8_reshaped = x_fp8.reshape(-1, orig_shape[-1])
@@ -52,22 +52,22 @@ def forward(
                 # calculate reference amax of output
                 with torch.no_grad():
                     ref_result = torch.addmm(b_fp8, x_fp8_reshaped, w_fp8.t())
-                    float8_amax_out.fill_(tensor_to_amax(ref_result))
+                    fp8_amax_y.fill_(tensor_to_amax(ref_result))
 
-            y_scale = amax_to_scale(float8_amax_out, torch.float8_e4m3fn)
+            y_scale = amax_to_scale(fp8_amax_y, torch.float8_e4m3fn)
             res_bits = addmm_float8(
-                b_fp8, x_fp8_reshaped, w_fp8.t(), float8_amax_out, y_scale, 
+                b_fp8, x_fp8_reshaped, w_fp8.t(), fp8_amax_y, y_scale, 
                 torch.float8_e4m3fn)
         else:
             if not is_fw_amax_initialized:
                 # calculate reference amax of output
                 with torch.no_grad():
                     ref_result = torch.mm(x_fp8_reshaped, w_fp8.t())
-                    float8_amax_out.fill_(tensor_to_amax(ref_result))
+                    fp8_amax_y.fill_(tensor_to_amax(ref_result))
 
-            y_scale = amax_to_scale(float8_amax_out, torch.float8_e4m3fn)
+            y_scale = amax_to_scale(fp8_amax_y, torch.float8_e4m3fn)
             res_bits = mm_float8(
-                x_fp8_reshaped, w_fp8.t(), float8_amax_out, y_scale, 
+                x_fp8_reshaped, w_fp8.t(), fp8_amax_y, y_scale, 
                 torch.float8_e4m3fn)
         res_bits = res_bits.reshape(*orig_shape[:-1], res_bits.shape[-1])
 
@@ -77,18 +77,18 @@ def forward(
 
     @staticmethod
     def backward(ctx, go):
-        x_fp8, w_fp8, b_fp8, float8_amax_dL_dX, float8_amax_dL_dW, \
-            float8_amax_dL_dY, bw_amax_initialized = \
+        x_fp8, w_fp8, b_fp8, fp8_amax_dL_dX, fp8_amax_dL_dW, \
+            fp8_amax_dL_dY, bw_amax_initialized = \
                 ctx.saved_tensors
 
         is_bw_amax_initialized = torch.any(bw_amax_initialized)
 
         if not isinstance(go, Float8Tensor):
             # TODO(future): switch to windowed delayed scaling
             if not is_bw_amax_initialized:
-                float8_amax_dL_dY.fill_(tensor_to_amax(go))
-            dL_dY_scale = amax_to_scale(float8_amax_dL_dY, torch.float8_e5m2)
-            float8_amax_dL_dY.fill_(tensor_to_amax(go))
+                fp8_amax_dL_dY.fill_(tensor_to_amax(go))
+            dL_dY_scale = amax_to_scale(fp8_amax_dL_dY, torch.float8_e5m2)
+            fp8_amax_dL_dY.fill_(tensor_to_amax(go))
             go_fp8 = Float8Tensor(
                 (go * dL_dY_scale).to(torch.float8_e5m2),
                 dL_dY_scale, go.dtype)
@@ -102,11 +102,11 @@ def backward(ctx, go):
             # calculate reference amax of output
             with torch.no_grad():
                 dL_dX_ref = torch.mm(go_fp8_reshaped, w_fp8)
-                float8_amax_dL_dX.fill_(tensor_to_amax(dL_dX_ref))
+                fp8_amax_dL_dX.fill_(tensor_to_amax(dL_dX_ref))
 
-        dL_dX_scale = amax_to_scale(float8_amax_dL_dX, torch.float8_e5m2)
+        dL_dX_scale = amax_to_scale(fp8_amax_dL_dX, torch.float8_e5m2)
         dL_dX_bits = mm_float8(
-            go_fp8_reshaped, w_fp8, float8_amax_dL_dX, dL_dX_scale, torch.float8_e5m2)
+            go_fp8_reshaped, w_fp8, fp8_amax_dL_dX, dL_dX_scale, torch.float8_e5m2)
         dL_dX_bits = dL_dX_bits.reshape(*go_fp8_orig_shape[:-1], dL_dX_bits.shape[-1])
         dL_dX_fp8 = Float8Tensor(dL_dX_bits, dL_dX_scale, go_fp8._orig_dtype)
 
@@ -117,11 +117,11 @@ def backward(ctx, go):
             # calculate reference amax of output
             with torch.no_grad():
                 dL_dW_ref = torch.mm(x_fp8_reshaped.t(), go_fp8_reshaped).t()
-                float8_amax_dL_dW.fill_(tensor_to_amax(dL_dW_ref))
+                fp8_amax_dL_dW.fill_(tensor_to_amax(dL_dW_ref))
 
-        dL_dW_scale = amax_to_scale(float8_amax_dL_dW, torch.float8_e5m2)
+        dL_dW_scale = amax_to_scale(fp8_amax_dL_dW, torch.float8_e5m2)
         dL_dW_bits = mm_float8(
-            x_fp8_reshaped.t(), go_fp8_reshaped, float8_amax_dL_dW, 
+            x_fp8_reshaped.t(), go_fp8_reshaped, fp8_amax_dL_dW, 
             dL_dW_scale, torch.float8_e5m2).t()
         dL_dW_fp8 = Float8Tensor(dL_dW_bits, dL_dW_scale, go_fp8._orig_dtype)
 
@@ -147,13 +147,13 @@ def __init__(self, *args, **kwargs):
         # scaling such as the mechanism described in
         # https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Mixed-precision-training-with-FP8,
         # or PTQ calibration.
-        self.register_buffer('float8_amax_in', torch.tensor(E4M3_MAX_POS))
-        self.register_buffer('float8_amax_weight', torch.tensor(E4M3_MAX_POS))
-        self.register_buffer('float8_amax_bias', torch.tensor(E4M3_MAX_POS))
-        self.register_buffer('float8_amax_out', torch.tensor(E4M3_MAX_POS))
-        self.register_buffer('float8_amax_dL_dX', torch.tensor(E5M2_MAX_POS))
-        self.register_buffer('float8_amax_dL_dW', torch.tensor(E5M2_MAX_POS))
-        self.register_buffer('float8_amax_dL_dY', torch.tensor(E5M2_MAX_POS))
+        self.register_buffer('fp8_amax_x', torch.tensor(E4M3_MAX_POS))
+        self.register_buffer('fp8_amax_w', torch.tensor(E4M3_MAX_POS))
+        self.register_buffer('fp8_amax_b', torch.tensor(E4M3_MAX_POS))
+        self.register_buffer('fp8_amax_y', torch.tensor(E4M3_MAX_POS))
+        self.register_buffer('fp8_amax_dL_dX', torch.tensor(E5M2_MAX_POS))
+        self.register_buffer('fp8_amax_dL_dW', torch.tensor(E5M2_MAX_POS))
+        self.register_buffer('fp8_amax_dL_dY', torch.tensor(E5M2_MAX_POS))
         self.register_buffer('fw_amax_initialized', torch.tensor([0], dtype=torch.uint8))
         self.register_buffer('bw_amax_initialized', torch.tensor([0], dtype=torch.uint8))
 
@@ -169,34 +169,34 @@ def forward(self, x):
 
             # TODO(future): switch to windowed delayed scaling
             if not is_fw_amax_initialized:
-                self.float8_amax_in.fill_(tensor_to_amax(x))
-            x_scale = amax_to_scale(self.float8_amax_in, torch.float8_e4m3fn)
-            self.float8_amax_in.fill_(tensor_to_amax(x))
+                self.fp8_amax_x.fill_(tensor_to_amax(x))
+            x_scale = amax_to_scale(self.fp8_amax_x, torch.float8_e4m3fn)
+            self.fp8_amax_x.fill_(tensor_to_amax(x))
 
             x_fp8 = Float8Tensor.to_float8(x, x_scale, torch.float8_e4m3fn)
         else:
             x_fp8 = x
 
         # TODO(future): switch to windowed delayed scaling
         if not is_fw_amax_initialized:
-            self.float8_amax_weight.fill_(tensor_to_amax(self.weight))
-        w_scale = amax_to_scale(self.float8_amax_weight, torch.float8_e4m3fn)
-        self.float8_amax_weight.fill_(tensor_to_amax(self.weight))
+            self.fp8_amax_w.fill_(tensor_to_amax(self.weight))
+        w_scale = amax_to_scale(self.fp8_amax_w, torch.float8_e4m3fn)
+        self.fp8_amax_w.fill_(tensor_to_amax(self.weight))
 
         w_fp8 = Float8Tensor.to_float8(self.weight, w_scale, torch.float8_e4m3fn)
         maybe_b_fp8 = None
         if self.bias is not None:
             # TODO(future): switch to windowed delayed scaling
             if not is_fw_amax_initialized:
-                self.float8_amax_bias.fill_(tensor_to_amax(self.bias))
-            b_scale = amax_to_scale(self.float8_amax_bias, torch.float8_e4m3fn)
-            self.float8_amax_bias.fill_(tensor_to_amax(self.bias))
+                self.fp8_amax_b.fill_(tensor_to_amax(self.bias))
+            b_scale = amax_to_scale(self.fp8_amax_b, torch.float8_e4m3fn)
+            self.fp8_amax_b.fill_(tensor_to_amax(self.bias))
 
             maybe_b_fp8 = Float8Tensor.to_float8(self.bias, b_scale, torch.float8_e4m3fn)
 
         y_fp8 = float8_linear.apply(
-            x_fp8, w_fp8, maybe_b_fp8, self.float8_amax_out, self.float8_amax_dL_dX,
-            self.float8_amax_dL_dW, self.float8_amax_dL_dY, self.fw_amax_initialized,
+            x_fp8, w_fp8, maybe_b_fp8, self.fp8_amax_y, self.fp8_amax_dL_dX,
+            self.fp8_amax_dL_dW, self.fp8_amax_dL_dY, self.fw_amax_initialized,
             self.bw_amax_initialized)
 
         if not is_fw_amax_initialized:

tests/test.py

@@ -88,15 +88,15 @@ def _test_linear_impl(self, x, m_ref):
 
         # verify all of the amax buffers got updated
         buffer_names = [
-            'float8_amax_in',
-            'float8_amax_weight',
-            'float8_amax_out',
-            'float8_amax_dL_dX',
-            'float8_amax_dL_dW',
-            'float8_amax_dL_dY',
+            'fp8_amax_x',
+            'fp8_amax_w',
+            'fp8_amax_y',
+            'fp8_amax_dL_dX',
+            'fp8_amax_dL_dW',
+            'fp8_amax_dL_dY',
         ]
         if m_ref.bias is not None:
-            buffer_names.append('float8_amax_bias')
+            buffer_names.append('fp8_amax_b')
         for buffer_name in buffer_names:
             buffer_value = getattr(m_fp8, buffer_name)
             for init_val in (E4M3_MAX_POS, E5M2_MAX_POS):