[BIONEMO-2334] Patch TE to fix Evo2 stop and go training

Signed-off-by: Bruno Alvisio <balvisio@nvidia.com>

Author: balvisio

Dockerfile

@@ -58,11 +58,21 @@ rm -rf /tmp/* /var/tmp/*
 EOF
 
 
-## BUMP TE as a solution to the issue https://github.com/NVIDIA/bionemo-framework/issues/422. Drop this when pytorch images ship the fixed commit.
- ARG TE_TAG=9d4e11eaa508383e35b510dc338e58b09c30be73
- RUN PIP_CONSTRAINT= NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi \
-    pip --disable-pip-version-check --no-cache-dir install \
-    git+https://github.com/NVIDIA/TransformerEngine.git@${TE_TAG}
+## BUMP and patch TE as a solution to the issues:
+## 1. https://github.com/NVIDIA/bionemo-framework/issues/422
+## 2. https://github.com/NVIDIA/bionemo-framework/issues/973
+## Drop this when pytorch images ship the fixed commit.
+
+# Set ARG
+ARG TE_TAG=9d4e11eaa508383e35b510dc338e58b09c30be73
+
+COPY ./patches/te.patch /tmp/te.patch
+RUN git clone --recurse-submodules https://github.com/NVIDIA/TransformerEngine.git /tmp/TransformerEngine && \
+    cd /tmp/TransformerEngine && \
+    git checkout --recurse-submodules ${TE_TAG} && \
+    patch -p1 < /tmp/te.patch && \
+    PIP_CONSTRAINT= NVTE_FRAMEWORK=pytorch NVTE_WITH_USERBUFFERS=1 MPI_HOME=/usr/local/mpi \
+    pip --disable-pip-version-check --no-cache-dir install .
 
 # Install AWS CLI based on architecture
 RUN if [ "$TARGETARCH" = "arm64" ]; then \

patches/te.patch

@@ -0,0 +1,29 @@
+diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
+index 2c3ccff1..a00d46cc 100644
+--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
++++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
+@@ -10,6 +10,7 @@
+ #include "extensions.h"
+ #include "pybind.h"
+ #include "transformer_engine/transformer_engine.h"
++#include "util.h"
+
+ namespace transformer_engine::pytorch {
+
+@@ -33,6 +34,16 @@ py::object quantize(const at::Tensor& tensor, py::handle quantizer, const py::ob
+     DType fake_te_type = GetTransformerEngineDType(fake_tensor_type);
+     std::tie(te_output, out) = my_quantizer->create_tensor(input_shape, fake_te_type);
+   } else {
++    if (my_quantizer->columnwise_usage && !non_tn_fp8_gemm_supported()) {
++      bool transpose_exists = !output.attr("_transpose_invalid").cast<bool>() && !output.attr("_transpose").is_none();
++      if (!transpose_exists) {
++        DType fake_te_type = GetTransformerEngineDType(fake_tensor_type);
++        py::object new_out;
++        std::tie(std::ignore, new_out) = my_quantizer->create_tensor(input_shape, fake_te_type);
++        output.attr("_transpose_invalid") = py::bool_(false);
++        output.attr("_transpose") = new_out.attr("_transpose");
++      }
++    }
+     out = output;
+     te_output = makeTransformerEngineTensor(output, quantizer);
+   }

sub-packages/bionemo-evo2/tests/bionemo/evo2/run/test_train.py

@@ -56,7 +56,7 @@ def small_training_cmd(path, max_steps, val_check, devices: int = 1, additional_
         "--model-size 1b_nv --num-layers 4 --hybrid-override-pattern SDH* --limit-val-batches 1 "
         "--no-activation-checkpointing --add-bias-output --create-tensorboard-logger --create-tflops-callback "
         f"--max-steps {max_steps} --warmup-steps 1 --val-check-interval {val_check} --limit-val-batches 1 "
-        f"--seq-length 8 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args}"
+        f"--seq-length 16 --hidden-dropout 0.1 --attention-dropout 0.1 {additional_args}"
     )
     return cmd
 
@@ -158,7 +158,6 @@ def test_train_evo2_stops(tmp_path):
             "--fp8",
             marks=[
                 pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
-                pytest.mark.xfail(reason="FP8 test currently broken - TODO: fix"),
             ],
             id="fp8",
         ),