Minor fixes

fast_llm/data/dataset/streaming.py

@@ -1,5 +1,6 @@
 import functools
 import json
+import logging
 import time
 import typing
 
@@ -14,6 +15,8 @@
 from fast_llm.data.document.token_data import TokenDataDocument
 from fast_llm.utils import Assert
 
+logger = logging.getLogger(__name__)
+
 
 @config_class()
 class RedisStreamingDocumentData(Config):

fast_llm/layers/ssm/gdn.py

@@ -227,7 +227,7 @@ def __init__(
             self._value_head_dim, lr_scale=self._lr_scale, peft=self._peft
         )
 
-        if _fast_gdn_available:
+        if _fast_gdn_available and distributed_config.use_cuda:
             self.chunk_gated_delta_rule = chunk_gated_delta_rule
         else:
             logger.warning(

fast_llm_external_models/apriel2/modeling_apriel2.py

@@ -2839,7 +2839,7 @@ def forward(
 
         # Reshape back to [batch, num_patches, text_hidden]
         image_features = image_features.squeeze(0).view(batch_size, num_patches_per_image, -1)
-        return image_features, (*all_hidden_states, hidden_states, image_features)
+        return image_features, (*all_hidden_states, hidden_states, image_features) if output_hidden_states else None
 
 
 class SimpleMLP(nn.Module):

fast_llm_external_models/tests/test_apriel2/test_equivalence.py

@@ -481,7 +481,7 @@ def test_batch_processing_behavior(self, model_pair):
         with torch.no_grad():
             # Batch processing
             batch_src = get_pixtral_vision_features(source, pixel_values)
-            batch_tgt, _ = target.get_image_features(pixel_values).view(-1, batch_src.shape[-1])
+            batch_tgt = target.get_image_features(pixel_values)[0].view(-1, batch_src.shape[-1])
 
             # Sequential processing
             singles_src = [get_pixtral_vision_features(source, pixel_values[i : i + 1]) for i in range(3)]

tests/models/test_streaming.py

@@ -132,7 +132,7 @@ def _run_model_streaming_configs(
             model_testing_config,
             None,
             updates={
-                ("data", "datasets"): {"training": {"port": port}},
+                ("data", "datasets"): {"training": {"port": port, "timeout": 1.0}},
                 ("training", "export"): {"format": model_testing_config.checkpoint_format.name, "interval": 1},
                 "callbacks": {
                     "streaming": {
@@ -143,6 +143,7 @@ def _run_model_streaming_configs(
                             "external_world_size": config.consumer_count,
                         },
                         "export": {"format": model_testing_config.checkpoint_format.name},
+                        "timeout": 1.0,
                     }
                 },
                 # Disable tensor logging.

tests/utils/distributed_configs.py

@@ -70,7 +70,7 @@ def get_config(relative: float = 0, absolute: float = 0, **kwargs) -> CompareCon
             if torch.cuda.is_available()
             else {
                 (None, "norm"): get_config(ignore_tensors=True),
-                (None, "word_embeddings_weight"): get_config(8e-2, 1e-4),
+                (None, "embeddings_weight"): get_config(8e-2, 1e-4),
             }
         ),
         (None, "bias"): get_config(2e-2, 1e-3) if torch.cuda.is_available() else get_config(2e-2, 2e-3),

tests/utils/model_configs.py

@@ -802,7 +802,7 @@ def update_and_add_testing_config(
     # note: tp is excluded because there is currently no gradient reductions implemented for tp norm in gdn.py (STP works though).
     # we should be using STP with this model, not TP!
     skip_tests=("sdp", "ms", TP_NO_STP),
-    requires_cuda=False,
+    requires_cuda=True,  # GDN available on CPU, but not in the converted model (also runs very slow).
 )
 
 _gdn_block = MODEL_CONFIGS["apriel2_gdn"].config_dict["model"]["base_model"]["decoder"]["block"]

tests/utils/redis.py

@@ -66,8 +66,6 @@ def producer_loop():
 
 @contextlib.contextmanager
 def fake_redis_server(config: RedisConfig):
-    # We search for free port as port from previous test can still be not free even after server shutdown
-
     # ----- Monkey-patch handler to suppress broken pipes -----
     orig_handle = fakeredis._tcp_server.TCPFakeRequestHandler.handle