[PyTorch] Debug weight matrix usages for dgrad GEMM (#1637)

Make sure that weight matrix has required usages for dgrad GEMM
Signed-off-by: Tim Moon <tmoon@nvidia.com>

transformer_engine/pytorch/module/layernorm_linear.py

@@ -327,7 +327,6 @@ class _LayerNormLinear(torch.autograd.Function):
                         ln_out.update_usage(rowwise_usage=False)
 
             # Weight with column-wise usage is needed for dgrad GEMM.
-            if inp.requires_grad:
             if isinstance(weightmat, QuantizedTensor):
                 weightmat.update_usage(columnwise_usage=True)
 

transformer_engine/pytorch/module/layernorm_mlp.py

@@ -415,7 +415,7 @@ class _LayerNormMLP(torch.autograd.Function):
         )
 
         # Weight with column-wise usage is needed for dgrad GEMM.
-        if is_grad_enabled and inp.requires_grad:
+        if is_grad_enabled:
             if isinstance(fc1_weight_final, QuantizedTensor):
                 fc1_weight_final.update_usage(columnwise_usage=True)
             if isinstance(fc2_weight_final, QuantizedTensor):