🤖 Hi @shralex, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified two distinct failure modes: a GPU collective operations conflict (due to a CUDA context pre-emption/corruption issue) and a transient TPU Pathways network timeout (infrastructure/environment flake).

1. GPU Collective Operations Failure (Codebase-Related / Environment)

🔍 What Failed

• Job/Matrix: gpu-unit and gpu-integration
• Failing Tests (Examples):
  • tests/unit/model_creation_utils_test.py::TestSetupDecodeStateFromNnx::test_returns_linen_train_state_and_annotations
  • tests/unit/multi_token_prediction_test.py::MultiTokenPredictionLayerTest::test_multi_token_prediction_layer_output
  • tests/integration/maxengine_test.py::MaxEngineTest::test_basic_decode_nnx
• Error: jax.errors.JaxRuntimeError: INTERNAL: NCCL operation failed: invalid argument / NCCL WARN Error: corrupted comm object detected

🪵 Error Details & Stack Trace

FAILED tests/unit/model_creation_utils_test.py::TestSetupDecodeStateFromNnx::test_returns_linen_train_state_and_annotations - jax.errors.JaxRuntimeError: INTERNAL: NCCL operation ncclAllReduce( send_buffer.opaque(), recv_buffer.opaque(), ToNcclCount(dtype, count), nccl_dtype, ToNcclReduction(reduction_kind), comm_, AsCudaStream(stream)) failed: invalid argument (run with NCCL_DEBUG=WARN for details). Last NCCL warning(error) log entry (may be unrelated) ''.

FAILED tests/unit/multi_token_prediction_test.py::MultiTokenPredictionLayerTest::test_multi_token_prediction_layer_output - jax.errors.JaxRuntimeError: INTERNAL: NCCL operation ncclAllGather( send_buffer.opaque(), recv_buffer.opaque(), ToNcclCount(dtype, count), nccl_dtype, comm_, AsCudaStream(stream)) failed: invalid argument (run with NCCL_DEBUG=WARN for details). Last NCCL warning(error) log entry (may be unrelated) ''.

linux-x86-a2-48-a100-4gpu-r2txm-runner-g6f9c-workflow:1125:1301 [1] external/nccl_archive/src/misc/argcheck.cc:39 NCCL WARN Error: corrupted comm object detected

💡 Root Cause Analysis & Context

Confidence: high (confirmed cause)

• The Conflict: In multi-GPU JAX environments, if TensorFlow or libraries importing TensorFlow (like Google's JAX-native quantization library Qwix or TFDS) are imported before JAX has explicitly initialized its devices, TensorFlow pre-emptively initializes the CUDA/NCCL context on the GPUs. When JAX later attempts to initialize its own NCCL communicators for multi-GPU collective operations (e.g., ncclAllReduce or ncclAllGather), a CUDA/NCCL state conflict occurs, causing NCCL to detect a "corrupted comm object" and abort with an invalid argument error.
• Why it triggered in this PR: This PR (deprecate-aqt-keep-deps) deprecates/removes the old AQT library and sets use_qwix_quantization to True by default in the base configuration (src/maxtext/configs/base.yml). Because of this, qwix is now imported early during test collection in several core modules (e.g., src/maxtext/layers/quantizations.py, src/maxtext/layers/moe.py, etc.). Since qwix includes LiteRT/TensorFlow serialization modules, its import pre-emptively loads the CUDA context.
• Why the existing workaround failed: tests/conftest.py has an early-initialization block that attempts to call jax.devices() early, but it restricts this check using strict environment variables:

  _has_gpu = (
      "cuda" in _jax_platforms
      or "gpu" in _jax_platforms
      ...
  )

  In the self-hosted GitHub Actions GPU runners, these environment variables are not set or are empty, so the early-initialization was bypassed, allowing Qwix/TensorFlow imports to initialize CUDA first and corrupt subsequent JAX collective runs.

🛠️ Recommended Fix

We can solve this robustly and permanently by unconditionally forcing JAX to initialize early inside tests/conftest.py with a simple try-except block. This ensures JAX always establishes its CUDA/NCCL context first, avoiding context pre-emption/corruption from any subsequent library imports (including TensorFlow, TFDS, or Qwix).

diff --git a/tests/conftest.py b/tests/conftest.py
index 99066cc..4fdfdfd 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -37,17 +37,9 @@ if not _absl_flags.FLAGS.is_parsed():
 import jax
 import os
 import importlib.util
 
-# Force early JAX initialization on GPU to prevent CUDA context conflicts with TensorFlow/PyTorch.
-# If JAX initialization is deferred, TensorFlow/PyTorch (imported during test collection)
-# might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail
-# with 'corrupted comm object detected'.
-def _initialize_jax_early():
-  _jax_platforms = os.getenv("JAX_PLATFORMS", "").lower()
-  _device_type = os.getenv("INPUTS_DEVICE_TYPE", "").lower()
-  _has_gpu = (
-      "cuda" in _jax_platforms
-      or "gpu" in _jax_platforms
-      or "cuda" in _device_type
-      or "gpu" in _device_type
-      or os.getenv("CUDA_VISIBLE_DEVICES") is not None
-      or os.getenv("NVIDIA_VISIBLE_DEVICES") is not None
-  )
-  if _has_gpu:
-    try:
-      _ = jax.devices()
-    except Exception:  # pylint: disable=broad-exception-caught
-      pass
+# Force early JAX initialization on all platforms to prevent CUDA context conflicts with TensorFlow/PyTorch/Qwix.
+# If JAX initialization is deferred, other libraries (imported during test collection)
+# might initialize CUDA first, causing JAX's subsequent NCCL communicator creation to fail
+# with 'corrupted comm object detected'.
+try:
+  _ = jax.devices()
+except Exception:  # pylint: disable=broad-exception-caught
+  pass

2. TPU Pathways Integration Failure (Infrastructure/Environment Flake)

🔍 What Failed

• Job/Matrix: maxtext_tpu_pathways_integration_tests
• Failing Component: Runner initialization / checkout stage
• Error: Error: 14 UNAVAILABLE: No connection established. Last error: Error: connect ECONNREFUSED 34.118.229.212:50051 followed by Error: backoff timeout

🪵 Error Details & Stack Trace

maxtext_tpu_pathways_integration_tests / run	UNKNOWN STEP	2026-06-13T06:13:49.7098315Z ##[warning]Retrying execution for ECONNREFUSED or UNAVAILABLE "Error: Error execing #!/bin/sh -l\n\ncd /__w/maxtext/maxtext && exec env ...
: Error: 14 UNAVAILABLE: No connection established. Last error: Error: connect ECONNREFUSED 34.118.229.212:50051".
maxtext_tpu_pathways_integration_tests / run	UNKNOWN STEP	2026-06-13T06:14:09.6992250Z ##[error]Error: backoff timeout

💡 Root Cause Analysis & Context

Confidence: high (confirmed cause)

This failure is a pure infrastructure/environment flake and is completely unrelated to the changes introduced in this PR.

• Details: The self-hosted runner failed to establish a network connection to the GKE/Pathways master endpoint 34.118.229.212:50051 during the runner's execution loop. This caused a connection refused error (ECONNREFUSED) which eventually hit the backoff timeout and aborted the job.
• Resolution: Re-triggering the TPU Pathways integration job when the cluster endpoints are healthy and reachable should resolve this failure.

🤖 Hi @sarunsingla11722, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

• Job/Matrix: cpu-unit tests / cpu-unit / run (which caused the overall CPU Tests (Matrix), GPU Tests (Matrix), and Pathways Integration checks to fail)
• Failing Tests:
  • tests/unit/managed_mldiagnostics_test.py::ManagedMLDiagnosticsTest::test_enabled_empty_region_passes_none
  • tests/unit/managed_mldiagnostics_test.py::ManagedMLDiagnosticsTest::test_enabled_populated_region_passes_region
  • tests/unit/managed_mldiagnostics_test.py::ManagedMLDiagnosticsTest::test_not_enabled_noop
• Error: AttributeError: type object 'ManagedMLDiagnostics' has no attribute 'mldiag'

🪵 Error Details & Stack Trace

AttributeError: type object 'ManagedMLDiagnostics' has no attribute 'mldiag'

💡 Root Cause Analysis & Context

Confidence: high

The test suite failed due to an AttributeError in the newly introduced unit tests in tests/unit/managed_mldiagnostics_test.py.

Specifically, the test file imports the ManagedMLDiagnostics class:

from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics

and then attempts to mock mldiag directly on the class ManagedMLDiagnostics:

with mock.patch.object(ManagedMLDiagnostics.mldiag, "machinelearning_run") as mock_run:

However, in src/maxtext/common/managed_mldiagnostics.py, mldiag is defined as a module-level variable:

mldiag, _ = mldiagnostics_modules()

It is not defined as a class attribute or property on the ManagedMLDiagnostics class.

🛠️ Recommended Fix

There are two ways to resolve this depending on the design preference:

Option A: Update the test to mock the module-level mldiag (Recommended)

This keeps the class namespace clean and avoids exposing global variables as class attributes purely for testing:

diff --git a/tests/unit/managed_mldiagnostics_test.py b/tests/unit/managed_mldiagnostics_test.py
index f3c4695..2c57f72 100644
--- a/tests/unit/managed_mldiagnostics_test.py
+++ b/tests/unit/managed_mldiagnostics_test.py
@@ -17,6 +17,7 @@
 import unittest
 from unittest import mock
 
+from maxtext.common import managed_mldiagnostics
 from maxtext.common.managed_mldiagnostics import ManagedMLDiagnostics
 import pytest
 
@@ -34,7 +35,7 @@ class ManagedMLDiagnosticsTest(unittest.TestCase):
     mock_config = mock.MagicMock()
     mock_config.managed_mldiagnostics = False
 
-    with mock.patch.object(ManagedMLDiagnostics.mldiag, "machinelearning_run") as mock_run:
+    with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run:
       ManagedMLDiagnostics(mock_config)
       mock_run.assert_not_called()
 
@@ -47,7 +48,7 @@ class ManagedMLDiagnosticsTest(unittest.TestCase):
     mock_config.managed_mldiagnostics_dir = "gs://test_dir"
     mock_config.get_keys.return_value = {"key1": "val1"}
 
-    with mock.patch.object(ManagedMLDiagnostics.mldiag, "machinelearning_run") as mock_run:
+    with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run:
       ManagedMLDiagnostics(mock_config)
       mock_run.assert_called_once_with(
           name="test_run",
@@ -66,7 +67,7 @@ class ManagedMLDiagnosticsTest(unittest.TestCase):
     mock_config.managed_mldiagnostics_dir = "gs://test_dir"
     mock_config.get_keys.return_value = {"key1": "val1"}
 
-    with mock.patch.object(ManagedMLDiagnostics.mldiag, "machinelearning_run") as mock_run:
+    with mock.patch.object(managed_mldiagnostics.mldiag, "machinelearning_run") as mock_run:
       ManagedMLDiagnostics(mock_config)
       mock_run.assert_called_once_with(
           name="test_run",

Option B: Expose mldiag as a class attribute on ManagedMLDiagnostics

If your design specifically intends for mldiag to be accessible on the ManagedMLDiagnostics class, add it as a class attribute:

diff --git a/src/maxtext/common/managed_mldiagnostics.py b/src/maxtext/common/managed_mldiagnostics.py
index 5f9b296..57c963a 100644
--- a/src/maxtext/common/managed_mldiagnostics.py
+++ b/src/maxtext/common/managed_mldiagnostics.py
@@ -31,6 +31,7 @@ class ManagedMLDiagnostics:
   """
 
   _instance = None  # Class attribute to hold the single instance
+  mldiag = mldiag   # Expose module-level mldiag as class attribute
 
   def __new__(cls, *args: Any, **kwargs: Any):
     """Overrides the instance creation method.

🤖 Hi @sarunsingla11722, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

🤖 CI Failure Investigation Report

I have analyzed the recent test failures in the CI pipeline and identified the following:

🔍 What Failed

• Job/Matrix: gpu-unit and gpu-integration
• Failing Tests (Cascade of NCCL communicator failures across unit and integration tests):
  • Unit test: tests/unit/model_creation_utils_test.py::TestSetupDecodeStateFromNnx::test_returns_linen_train_state_and_annotations
  • Unit test: tests/unit/multi_token_prediction_test.py::MultiTokenPredictionLayerTest::test_multi_token_prediction_layer_output
  • Integration test: tests/integration/train_tests.py::TrainTests::test_gpu_hf_input_pipeline
• Error: jax.errors.JaxRuntimeError: INTERNAL: NCCL operation ncclAllReduce / ncclAllGather failed: invalid argument (run with NCCL_DEBUG=WARN for details). Last NCCL warning(error) log entry (may be unrelated) ''. corrupted comm object detected.

🪵 Error Details & Stack Trace

E         jax.errors.JaxRuntimeError: INTERNAL: NCCL operation ncclAllReduce( send_buffer.opaque(), recv_buffer.opaque(), ToNcclCount(dtype, count), nccl_dtype, ToNcclReduction(reduction_kind), comm_, AsCudaStream(stream)) failed: invalid argument (run with NCCL_DEBUG=WARN for details). Last NCCL warning(error) log entry (may be unrelated) ''.
...
E         1125:1303 [2] external/nccl_archive/src/misc/argcheck.cc:39 NCCL WARN Error: corrupted comm object detected

💡 Root Cause Analysis & Context

Confidence: Moderate

1. Verdict of the Isolated Debug Experiment

This PR was constructed specifically to run a clean debug experiment to isolate the root cause of the GPU/NCCL test failures:

• The Experiment: Keep the aqtp dependency in lockfiles/requirements but include 100% of the layer/model code changes (decoupling type hints, pruning in_serve_mode, transitioning to Qwix/FP8).
• The Verdict (Scenario B): The GPU tests on this branch STILL FAIL with the exact same NCCL corrupted communicator error.
• Key Finding: We can completely rule out the dependency removal or lockfile modifications as the culprit. The issue is triggered by either the layer/model code changes or a system-level GPU runner environment/NCCL network configuration mismatch.

2. Deep Dive: Logical Bug vs. Environment/Runner Flake

Standard tests like TestSetupDecodeStateFromNnx (which construct a minimal LLaMA2 model with quantization set to "" / completely disabled) also fail with the identical NCCL error. Since no MoE layers or Qwix/FP8 quantization operations are ever executed in these standard tests, it strongly suggests a system-level runner environment issue rather than a localized bug inside the new layer logic.

A likely environment-level culprit is the removal of the network-binding environment variable NCCL_SOCKET_IFNAME=lo in .github/workflows/run_tests_against_package.yml. On self-hosted GHA runner virtual machines running multi-GPU tests, omitting NCCL_SOCKET_IFNAME=lo forces NCCL to attempt to bind to external interface sockets (like eth0) rather than the local loopback interface, leading to TCP timeouts/connection issues that corrupt the NCCL communicator object and crash the entire test suite.