[Common] Remove CheckTensor if the workspace is empty in cast_transpose_fused (#931)

* rm tensor check if the workspace is empty
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* add trust_remote=true for load_dataset() in the mnist test
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

examples/jax/mnist/test_single_gpu_mnist.py

@@ -128,7 +128,7 @@ def eval_model(state, test_ds, batch_size, var_collect):
 
 def get_datasets():
     """Load MNIST train and test datasets into memory."""
-    train_ds = load_dataset("mnist", split="train")
+    train_ds = load_dataset("mnist", split="train", trust_remote_code=True)
     train_ds.set_format(type="np")
     batch_size = train_ds["image"].shape[0]
     shape = (batch_size, IMAGE_H, IMAGE_W, IMAGE_C)
@@ -136,7 +136,7 @@ def get_datasets():
         "image": train_ds["image"].astype(np.float32).reshape(shape) / 255.0,
         "label": train_ds["label"],
     }
-    test_ds = load_dataset("mnist", split="test")
+    test_ds = load_dataset("mnist", split="test", trust_remote_code=True)
     test_ds.set_format(type="np")
     batch_size = test_ds["image"].shape[0]
     shape = (batch_size, IMAGE_H, IMAGE_W, IMAGE_C)

transformer_engine/common/transpose/cast_transpose_fusion.cu

@@ -478,9 +478,13 @@ template <bool IS_DBIAS, bool IS_DACT, typename ComputeType, typename ParamOP,
 void cast_transpose_fused(const Tensor &input, const Tensor &act_input, Tensor *cast_output,
                           Tensor *transposed_output, Tensor *dbias, Tensor *workspace,
                           cudaStream_t stream) {
-  CheckInputTensor(input, "cast_transpose_fused_input");
-  CheckOutputTensor(*cast_output, "cast_output");
-  CheckOutputTensor(*transposed_output, "transposed_output");
+  if (workspace->data.dptr != nullptr) {
+    CheckInputTensor(input, "cast_transpose_fused_input");
+    CheckOutputTensor(*cast_output, "cast_output");
+    CheckOutputTensor(*transposed_output, "transposed_output");
+    if constexpr (IS_DBIAS) CheckOutputTensor(*dbias, "dbias");
+    if constexpr (IS_DACT) CheckInputTensor(act_input, "act_input");
+  }
 
   NVTE_CHECK(input.data.shape.size() == 2, "Input must have 2 dimensions.");
   NVTE_CHECK(cast_output->data.shape.size() == 2, "C output must have 2 dimensions.");
@@ -501,12 +505,10 @@ void cast_transpose_fused(const Tensor &input, const Tensor &act_input, Tensor *
              "C and T outputs need to share scale tensor.");
 
   if constexpr (IS_DBIAS) {
-    CheckOutputTensor(*dbias, "dbias");
     NVTE_CHECK(dbias->data.dtype == input.data.dtype, "DBias must have the same type as input.");
     NVTE_CHECK(dbias->data.shape == std::vector<size_t>{row_length}, "Wrong shape of DBias.");
   }
   if constexpr (IS_DACT) {
-    CheckInputTensor(act_input, "act_input");
     NVTE_CHECK(input.data.dtype == act_input.data.dtype, "Types of both inputs must match.");
     NVTE_CHECK(input.data.shape == act_input.data.shape, "Shapes of both inputs must match.");
   }