[tx] General implementation of trainable Hyper Connections

skyrl-tx/tx/layers/connectors.py

@@ -0,0 +1,104 @@
+from flax import nnx
+import jax
+from jax import numpy as jnp
+
+from tx.layers.util import Param
+from tx.layers.layernorm import RMSNorm
+
+
+class Connector(nnx.Module):
+    """
+    Implementation of Manifold constrained HyperConnections (https://arxiv.org/pdf/2512.24880)
+
+    Weights initialized with identity mapping; Default behaviour equates to residual networks.
+    """
+
+    def __init__(
+        self,
+        hidden_dim: int,
+        expansion_rate: int,
+        *,
+        trainable: bool = False,
+        sinkhorn_iters: int = 20,
+        eps: float = 1e-5,
+        dtype: jnp.dtype,
+        rngs: nnx.Rngs,
+    ) -> None:
+        self.hidden_dim = hidden_dim
+        self.expansion_rate = expansion_rate
+        self.trainable = trainable
+        self.sinkhorn_iters = sinkhorn_iters
+        self.eps = eps
+        n = expansion_rate
+        C = hidden_dim
+
+        self.norm = RMSNorm(hidden_dim, eps=eps, dtype=dtype, rngs=rngs)
+
+        self.phi_pre = Param(n * C, n, dtype=dtype, kernel_init=nnx.initializers.normal(stddev=0.02), rngs=rngs)
+        self.phi_post = Param(n * C, n, dtype=dtype, kernel_init=nnx.initializers.normal(stddev=0.02), rngs=rngs)
+        self.phi_res = Param(n * C, n * n, dtype=dtype, kernel_init=nnx.initializers.normal(stddev=0.02), rngs=rngs)
+
+        # Initialize biases for identity-like behavior:
+        # H_pre = 1/n (uniform aggregation), H_post = 1 (full distribution), M = I (identity mixing)
+
+        # H_pre = sigmoid(b_pre) = 1/n  =>  b_pre = logit(1/n)
+        target_h_pre = jnp.array(1.0 / n, dtype=dtype)
+        clamped = jnp.clip(target_h_pre, 1e-6, 1.0 - 1e-6)
+        logit_1_over_n = jnp.log(clamped) - jnp.log(1.0 - clamped)
+        self.b_pre = nnx.Param(jnp.full((n,), logit_1_over_n, dtype=dtype))
+
+        # H_post = 2 * sigmoid(b_post) = 1  =>  b_post = 0
+        self.b_post = nnx.Param(jnp.zeros((n,), dtype=dtype))
+
+        # M = sinkhorn(exp(b_res)) = I  =>  b_res = large diagonal matrix
+        self.b_res = nnx.Param(10.0 * jnp.eye(n, dtype=dtype))
+
+        self.alpha_pre = nnx.Param(jnp.array(0.0, dtype=dtype))
+        self.alpha_post = nnx.Param(jnp.array(0.0, dtype=dtype))
+        self.alpha_res = nnx.Param(jnp.array(0.0, dtype=dtype))
+
+    def _sinkhorn_knopp(self, M: jax.Array) -> jax.Array:
+        M = jnp.exp(M)
+        for _ in range(self.sinkhorn_iters):
+            M = M / (M.sum(axis=-1, keepdims=True) + self.eps)
+            M = M / (M.sum(axis=-2, keepdims=True) + self.eps)
+        return M
+
+    def _get_params(self):
+        """Get all connector params, with stop_gradient applied if not trainable."""
+        sg = (lambda x: x) if self.trainable else jax.lax.stop_gradient
+        return (
+            sg(self.alpha_pre[...]), sg(self.alpha_post[...]), sg(self.alpha_res[...]),
+            sg(self.phi_pre[...]), sg(self.phi_post[...]), sg(self.phi_res[...]),
+            sg(self.b_pre[...]), sg(self.b_post[...]), sg(self.b_res[...]),
+            sg(self.norm.weight[...]),
+        )
+
+    def pre(self, x: jax.Array) -> jax.Array:
+        *batch_dims, n, C = x.shape
+
+        x_flat = x.reshape(*batch_dims, n * C)
+        rms = jnp.sqrt(jnp.mean(x_flat * x_flat, axis=-1, keepdims=True) + self.eps)
+        x_norm = x_flat / rms
+
+        (alpha_pre, alpha_post, alpha_res, phi_pre, phi_post, phi_res,
+         b_pre, b_post, b_res, norm_weight) = self._get_params()
+
+        tilde_H_pre = alpha_pre * (x_norm @ phi_pre) + b_pre
+        tilde_H_post = alpha_post * (x_norm @ phi_post) + b_post
+        tilde_H_res = alpha_res * (x_norm @ phi_res).reshape(*batch_dims, n, n) + b_res
+
+        H_pre = jax.nn.sigmoid(tilde_H_pre)
+        self._H_post = 2.0 * jax.nn.sigmoid(tilde_H_post)
+        self._M = self._sinkhorn_knopp(tilde_H_res)
+
+        x_agg = jnp.einsum("...i,...ic->...c", H_pre, x)
+        rms_norm = jnp.sqrt(jnp.mean(x_agg**2, axis=-1, keepdims=True) + self.norm.eps)
+        x_normed = norm_weight * x_agg / rms_norm
+
+        return x_normed
+
+    def post(self, residual: jax.Array, output: jax.Array) -> jax.Array:
+        y_dist = self._H_post[..., None] * output[..., None, :]
+        x_mixed = jnp.einsum("...ij,...jc->...ic", self._M, residual)
+        return x_mixed + y_dist

skyrl-tx/tx/layers/layernorm.py

@@ -16,7 +16,7 @@ class RMSNorm(nnx.Module):
     def __init__(self, size: int, *, eps: float = 1e-6, dtype: jnp.dtype, rngs: nnx.Rngs) -> None:
         self.eps = eps
         self.weight = Param(
-            size, dtype=dtype, kernel_init=nnx.with_partitioning(nnx.initializers.normal(), jax.P(None)), rngs=rngs
+            size, dtype=dtype, kernel_init=nnx.with_partitioning(nnx.initializers.ones_init(), jax.P(None)), rngs=rngs
         )
 
     def __call__(self, x: jax.Array) -> jax.Array:

skyrl-tx/tx/models/deepseekv3.py

@@ -7,6 +7,7 @@
 from tx.layers.rotary_embedding import get_rope
 from tx.layers.util import Param, prepare_routing, shard_map_ep
 from tx.layers.layernorm import RMSNorm
+from tx.layers.connectors import Connector
 from tx.models.configs import DeepseekV3Config
 from tx.models.types import CausalLMOutput, ModelForCausalLM, ModelOutput
 from tx.utils.generator import GeneratorMixin, KVCache
@@ -417,17 +418,28 @@ def __call__(
 
 class DeepseekV3DecoderLayer(nnx.Module):
 
-    def __init__(self, config: DeepseekV3Config, layer_idx: int, *, dtype: jnp.dtype, rngs: nnx.Rngs) -> None:
-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype, rngs=rngs)
-        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, dtype=dtype, rngs=rngs)
+    def __init__(
+        self,
+        config: DeepseekV3Config,
+        layer_idx: int,
+        *,
+        dtype: jnp.dtype,
+        rngs: nnx.Rngs,
+        expansion_rate: int = 1,
+    ) -> None:
         self.self_attn = DeepseekV3Attention(config, dtype=dtype, rngs=rngs)
+        self.layer_idx = layer_idx
+        self.num_layers = config.num_hidden_layers
+        self.expansion_rate = expansion_rate
 
-        # Use dense MLP for initial layers, MoE for the rest
         if layer_idx >= config.first_k_dense_replace:
             self.mlp = DeepseekV3MoE(config, dtype=dtype, rngs=rngs)
         else:
             self.mlp = DeepseekV3MLP(config, dtype=dtype, rngs=rngs)
 
+        self.attn_connector = Connector(config.hidden_size, expansion_rate, dtype=dtype, rngs=rngs)
+        self.mlp_connector = Connector(config.hidden_size, expansion_rate, dtype=dtype, rngs=rngs)
+
     def __call__(
         self,
         hidden_states: jax.Array,
@@ -437,21 +449,28 @@ def __call__(
         adapter_indices: jax.Array | None = None,
         kv_cache: tuple[jax.Array, jax.Array] | None = None,
     ) -> tuple[jax.Array, tuple[jax.Array, jax.Array]]:
+        n = self.expansion_rate
+        if self.layer_idx == 0:
+            hidden_states = jnp.repeat(hidden_states[..., None, :], n, axis=-2)
+
         residual = hidden_states
-        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states = self.attn_connector.pre(hidden_states)
         hidden_states, updated_cache = self.self_attn(
             hidden_states,
             attention_mask=attention_mask,
             positions=positions,
             adapter_indices=adapter_indices,
             kv_cache=kv_cache,
         )
-        hidden_states = residual + hidden_states
+        hidden_states = self.attn_connector.post(residual, hidden_states)
 
         residual = hidden_states
-        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp_connector.pre(hidden_states)
         mlp_output = self.mlp(hidden_states, adapter_indices=adapter_indices)
-        hidden_states = residual + mlp_output
+        hidden_states = self.mlp_connector.post(residual, mlp_output)
+
+        if self.layer_idx == self.num_layers - 1:
+            hidden_states = hidden_states.sum(axis=-2)
 
         return hidden_states, updated_cache
 
@@ -500,7 +519,7 @@ def __call__(
 
         for layer_idx, layer in enumerate(self.layers):
             if output_hidden_states:
-                all_hidden_states.append(hidden_states)
+                all_hidden_states.append(hidden_states.squeeze())
 
             hidden_states, (k, v) = layer(
                 hidden_states,

skyrl-tx/tx/utils/models.py

@@ -113,6 +113,9 @@ def load_safetensors(
         # Skip LoRA parameters if requested
         if skip_lora and ("lora_A" in path or "lora_B" in path or "lora_scaling" in path or "lora_ranks" in path):
             continue
+        # Skip connector parameters
+        if any("connector" in str(p) for p in path):
+            continue
         if "experts" in path:
             tensors[key] = np.stack(
                 [tensors[get_expert_key(path, i)].T for i in range(config.get_num_experts())], axis=0