Release common-recipe model and update model card

Publish the frozen GPU-trained shared-prompt model with verified files and synchronized documentation.

README.md

@@ -51,27 +51,27 @@ The model was trained with data from the
 
 ## Evaluation
 
-Results on the held-out multi-domain test set:
+Results on the internal multi-domain test set:
 
 | Metric | Result |
 | :--- | :---: |
-| Test accuracy (%) | 86.82 |
+| Test accuracy (%) | 86.99 |
 | Scoring Spearman | 0.7264 |
-| Coherence accuracy | 75.59% |
-| Commonsense accuracy | 97.39% |
-| Empathy accuracy | 92.72% |
-| Multicultural accuracy | 74.99% |
+| Coherence accuracy | 76.32% |
+| Commonsense accuracy | 97.33% |
+| Empathy accuracy | 93.30% |
+| Multicultural accuracy | 74.31% |
 
 ## Hugging Face Models
 
 | Model | Base reward model | Test accuracy (%) | Scoring Spearman |
 | :--- | :--- | :---: | :---: |
-| [**`multi-domain-rm-fsfairx-gemma-2-9b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-gemma-2-9b-it) | [sfairXC/FsfairX-Gemma2-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-Gemma2-RM-v0.1) | **88.01** | 0.7346 |
-| [**`multi-domain-rm-skywork-qwen-3-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-skywork-qwen-3-8b-it) | [Skywork/Skywork-Reward-V2-Qwen3-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Qwen3-8B) | **87.82** | 0.7156 |
-| [**`multi-domain-rm-fsfairx-llama-3-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-llama-3-8b-it) | [sfairXC/FsfairX-LLaMA3-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1) | **86.86** | 0.7108 |
-| [**`multi-domain-rm-skywork-llama-3.1-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it) | [Skywork/Skywork-Reward-V2-Llama-3.1-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B) | **86.82** | 0.7264 |
-| [**`multi-domain-rm-mistral-7b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-mistral-7b-it) | [weqweasdas/RM-Mistral-7B](https://huggingface.co/weqweasdas/RM-Mistral-7B) | **84.41** | 0.6710 |
-| [**`multi-domain-rm-qwen-3-nemotron-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-qwen-3-nemotron-8b-it) | [nvidia/Qwen3-Nemotron-8B-BRRM](https://huggingface.co/nvidia/Qwen3-Nemotron-8B-BRRM) | **83.65** | 0.6704 |
+| [**`multi-domain-rm-fsfairx-gemma-2-9b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-gemma-2-9b-it) | [sfairXC/FsfairX-Gemma2-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-Gemma2-RM-v0.1) | **88.80** | 0.7346 |
+| [**`multi-domain-rm-skywork-qwen-3-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-skywork-qwen-3-8b-it) | [Skywork/Skywork-Reward-V2-Qwen3-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Qwen3-8B) | **88.08** | 0.7156 |
+| [**`multi-domain-rm-fsfairx-llama-3-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-llama-3-8b-it) | [sfairXC/FsfairX-LLaMA3-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1) | **87.75** | 0.7108 |
+| [**`multi-domain-rm-skywork-llama-3.1-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it) | [Skywork/Skywork-Reward-V2-Llama-3.1-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B) | **86.99** | 0.7264 |
+| [**`multi-domain-rm-mistral-7b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-mistral-7b-it) | [weqweasdas/RM-Mistral-7B](https://huggingface.co/weqweasdas/RM-Mistral-7B) | **85.25** | 0.6710 |
+| [**`multi-domain-rm-qwen-3-nemotron-8b-it`**](https://huggingface.co/mario-rc/multi-domain-rm-qwen-3-nemotron-8b-it) | [nvidia/Qwen3-Nemotron-8B-BRRM](https://huggingface.co/nvidia/Qwen3-Nemotron-8B-BRRM) | **84.35** | 0.6704 |
 
 ## Usage
 
@@ -86,7 +86,7 @@ repo_id = "mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it"
 tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
 model = AutoModel.from_pretrained(
     repo_id,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     device_map="auto",
     trust_remote_code=True,
 ).eval()
@@ -98,41 +98,48 @@ chosen = prompt + [{
 }]
 rejected = prompt + [{"role": "assistant", "content": "Tell them to ignore it."}]
 
-prompt_ids = tokenizer.apply_chat_template(
+prompt_inputs = tokenizer.apply_chat_template(
     prompt,
     tokenize=True,
     add_generation_prompt=True,
     return_tensors="pt",
+    return_dict=True,
 ).to(model.device)
-chosen_ids = tokenizer.apply_chat_template(
+chosen_inputs = tokenizer.apply_chat_template(
     chosen,
     tokenize=True,
     add_generation_prompt=False,
     return_tensors="pt",
+    return_dict=True,
 ).to(model.device)
-rejected_ids = tokenizer.apply_chat_template(
+rejected_inputs = tokenizer.apply_chat_template(
     rejected,
     tokenize=True,
     add_generation_prompt=False,
     return_tensors="pt",
+    return_dict=True,
 ).to(model.device)
 
 with torch.inference_mode():
-    gate = model.compute_gating(input_ids=prompt_ids)
+    gate = model.compute_gating(
+        input_ids=prompt_inputs["input_ids"],
+        attention_mask=prompt_inputs["attention_mask"],
+    )
     chosen_score = model(
-        input_ids=chosen_ids,
+        input_ids=chosen_inputs["input_ids"],
+        attention_mask=chosen_inputs["attention_mask"],
         gating_output_override=gate,
     ).score
     rejected_score = model(
-        input_ids=rejected_ids,
+        input_ids=rejected_inputs["input_ids"],
+        attention_mask=rejected_inputs["attention_mask"],
         gating_output_override=gate,
     ).score
 
 print({"chosen": chosen_score.item(), "rejected": rejected_score.item()})
 ```
 
-For padded batches, also pass the matching `attention_mask`. Scores are intended for comparison
-within a prompt; they are not calibrated probabilities or universal utility values.
+Pass the tokenizer's `input_ids` tensor and matching `attention_mask` to the model. Reuse the same prompt-derived gate for both candidates. Scores are intended for comparison within a prompt; they are not calibrated probabilities or universal utility values.
 
 ## Limitations
 
@@ -141,6 +148,8 @@ and should be calibrated for each downstream use case. Performance can vary by l
 distribution. The model inherits limitations and biases from its base model and training data and
 should not be used as the sole decision-maker in high-impact settings.
 
+The internal test was examined during development, and a source audit identified some train–test prompt overlap. These results are not an independent confirmation of generalization.
+
 ## Credits
 
 This model is based on the ArmoRM/RLHFlow reward-modeling approach and adapts it to custom

config.json

@@ -42,7 +42,7 @@
   "gating_dropout": 0.05,
   "gating_hidden_dim": 64,
   "gating_learnable_logit_scale": false,
-  "gating_logit_scale": 2.0,
+  "gating_logit_scale": 4.0,
   "gating_n_hidden": 1,
   "gating_temperature": 2.0,
   "head_dim": 128,

model-00001-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:86ca401765f707e08ad005924b409d4db9d332a95001c903a56bd1df40368657
+oid sha256:5aa0b32b04e4bd8ed6bd46cb09cec0212a8b532c2b7b0900d1bd49f3e6d82f81
 size 3903440936

modeling_custom.py

@@ -18,6 +18,8 @@ class GatingNetwork(nn.Module):
                  learnable_logit_scale: bool = False,
                  active_attribute_indices: Optional[List[int]] = None):
         super().__init__()
+        if temperature <= 0:
+            raise ValueError("Temperature must be positive.")
         self.temperature = temperature
         self.logit_scale = nn.Parameter(
             torch.ones(1) * logit_scale, requires_grad=learnable_logit_scale
@@ -25,12 +27,18 @@ class GatingNetwork(nn.Module):
         self.dropout_prob = dropout
         active_mask = torch.ones(out_features, dtype=torch.bool)
         if active_attribute_indices is not None:
-            if not active_attribute_indices:
-                raise ValueError("At least one attribute must be active.")
+            if (not isinstance(active_attribute_indices, (list, tuple))
+                    or not active_attribute_indices
+                    or any(type(index) is not int for index in active_attribute_indices)
+                    or len(set(active_attribute_indices)) != len(active_attribute_indices)):
+                raise ValueError("Active attribute indices must be unique integers in a nonempty sequence.")
             if min(active_attribute_indices) < 0 or max(active_attribute_indices) >= out_features:
                 raise ValueError("active_attribute_indices contains an out-of-range index.")
             active_mask.zero_()
             active_mask[list(active_attribute_indices)] = True
+        self.active_attribute_indices = tuple(
+            range(out_features) if active_attribute_indices is None else active_attribute_indices
+        )
         # Derived from packaged config; omit from state_dict for legacy compatibility.
         self.register_buffer("active_attribute_mask", active_mask, persistent=False)
         layers = []
@@ -40,6 +48,17 @@ class GatingNetwork(nn.Module):
         layers.append(nn.Linear(in_features, out_features, bias=bias))
         self.layers = nn.ModuleList(layers)
 
+    @torch.no_grad()
+    def reset_active_attribute_mask(self):
+        """Restore config-derived state after a low-memory/meta-device load.
+
+        Transformers can materialize non-persistent buffers with empty_like;
+        their constructor values therefore are not sufficient initialization.
+        This method changes no trained parameter or persistent checkpoint key.
+        """
+        self.active_attribute_mask.zero_()
+        self.active_attribute_mask[list(self.active_attribute_indices)] = True
+
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         for i, layer in enumerate(self.layers):
             x = layer(x)
@@ -68,24 +87,33 @@ class RewardModelWithGating(PreTrainedModel):
     config_class = AutoConfig
     base_model_prefix = "model"
 
+    def _init_weights(self, module):
+        if isinstance(module, GatingNetwork):
+            # The non-persistent mask is absent from saved weights by design.
+            # Initialize it when the loader materializes missing/derived state,
+            # without reinitializing the gating parameters it already loaded.
+            module.reset_active_attribute_mask()
+        else:
+            super()._init_weights(module)
+
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
         self.model = AutoModel.from_config(config)
         config_dict = config.to_dict()
-        
+
         # Default objective count for this project.
         self.num_objectives = config_dict.get("num_objectives", 23)
-        
+
         self.regression_layer = nn.Linear(config.hidden_size, self.num_objectives, bias=False)
         self.post_init()
-        
+
         # Avoid torch.eye to keep compatibility with BF16 training setups.
         I = torch.zeros(self.num_objectives, self.num_objectives)
         I[range(self.num_objectives), range(self.num_objectives)] = 1.
         self.reward_transform_matrix = nn.Parameter(I)
         self.reward_transform_matrix.requires_grad = False
-        
+
         self.gating = GatingNetwork(config.hidden_size, self.num_objectives,
                                     temperature=config_dict.get("gating_temperature", 10),
                                     logit_scale=config_dict.get("gating_logit_scale", 1.0),
@@ -159,10 +187,10 @@ class RewardModelWithGating(PreTrainedModel):
             batch_size = inputs_embeds.shape[0]
         else:
             raise ValueError("Either input_ids or inputs_embeds must be provided.")
-            
+
         if self.config.pad_token_id is None and batch_size != 1:
             raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
-            
+
         if self.config.pad_token_id is None:
             sequence_lengths = -1
         else:
@@ -173,11 +201,11 @@ class RewardModelWithGating(PreTrainedModel):
                 sequence_lengths = sequence_lengths.to(tokens_hidden_states.device)
             else:
                 sequence_lengths = -1
-                
+
         dummy_iterator = torch.arange(batch_size, device=tokens_hidden_states.device)
         hidden_states = tokens_hidden_states[dummy_iterator, sequence_lengths]
         assert hidden_states.shape == (batch_size, self.config.hidden_size)
-        
+
         rewards = self.regression_layer(hidden_states)
         prompt_embedding = None
         if gating_output_override is not None:
@@ -189,6 +217,11 @@ class RewardModelWithGating(PreTrainedModel):
                     f"gating_output_override shape {tuple(gating_output.shape)} does not "
                     f"match rewards shape {tuple(rewards.shape)}"
                 )
+        elif getattr(self.config, "shared_prompt_gating", False):
+            raise ValueError(
+                "This shared-prompt-gating checkpoint requires gating_output_override. "
+                "Compute one prompt-only gate with compute_gating() and reuse it for all candidates."
+            )
         else:
             if input_ids is None:
                 raise ValueError("input_ids is required to compute gating token positions.")
@@ -200,7 +233,7 @@ class RewardModelWithGating(PreTrainedModel):
             gating_output = self.gating(prompt_embedding)
         rewards_adjusted = rewards @ self.reward_transform_matrix
         score = torch.sum(gating_output * rewards_adjusted, dim=1)
-        
+
         return CustomOutput(
             rewards=rewards,
             hidden_state=hidden_states,
@@ -212,4 +245,4 @@ class RewardModelWithGating(PreTrainedModel):
 
 
 # Backward compatibility alias for existing imports/checkpoints.
-LlamaForRewardModelWithGating = RewardModelWithGating
+LlamaForRewardModelWithGating = RewardModelWithGating

requirements.txt

@@ -1,4 +1,4 @@
 torch>=2.6
-transformers>=4.51
+transformers==5.3.0
 safetensors>=0.5
 accelerate>=1.0

utils.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import importlib.util
 from typing import Optional, Sequence
 
 import torch
@@ -18,6 +19,144 @@ def _requires_remote_code(model_path: str) -> bool:
     return "qwen3" in model_path_l
 
 
+def _attention_implementation(device: str) -> str | None:
+    """Use FlashAttention on CUDA when installed; otherwise use Transformers defaults."""
+    if str(device).startswith("cuda") and importlib.util.find_spec("flash_attn") is not None:
+        return "flash_attention_2"
+    return None
+
+
+def _stable_int64_id(value) -> int:
+    """Return a deterministic non-negative signed-int64 identifier."""
+    import hashlib
+
+    if not isinstance(value, str):
+        value = json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
+    digest = hashlib.blake2b(value.encode("utf-8"), digest_size=8).digest()
+    return int.from_bytes(digest, "big", signed=False) & ((1 << 63) - 1)
+
+
+def debiasing_checkpoint_suffix(debiasing_dims, corr_threshold: float) -> str:
+    """Encode reward-transform settings in a stable checkpoint suffix."""
+    dims = sorted({int(dimension) for dimension in (debiasing_dims or ()) if int(dimension) >= 0})
+    if not dims:
+        return "_dbnone"
+    threshold = format(float(corr_threshold), ".12g").replace("-", "m").replace(".", "p")
+    dimension_text = "-".join(map(str, dims))
+    return f"_db{dimension_text}_ct{threshold}"
+
+
+def validate_shared_routing_config(routing_config) -> dict:
+    """Validate the metadata contract required by packaged shared-gate checkpoints."""
+    if not isinstance(routing_config, dict):
+        raise ValueError("Stage 2 checkpoint is missing its training_config mapping.")
+    if routing_config.get("format_version") != 2 or not routing_config.get("shared_prompt_gating", False):
+        raise ValueError(
+            "Stage 2 checkpoint must declare format_version=2 and "
+            "shared_prompt_gating=true; legacy checkpoints are not packageable."
+        )
+    return routing_config
+
+
+def score_shared_gate_candidates(
+    candidate_attribute_rewards: torch.Tensor,
+    shared_gate_weights: torch.Tensor,
+) -> torch.Tensor:
+    """Combine per-candidate attribute rewards with one gate per prompt pair."""
+    if candidate_attribute_rewards.ndim != 3 or shared_gate_weights.ndim != 2:
+        raise ValueError(
+            "Shared-gate scoring expects candidate rewards shaped "
+            "[pairs, candidates, attributes] and gate weights shaped "
+            "[pairs, attributes]."
+        )
+    if (
+        candidate_attribute_rewards.shape[0] != shared_gate_weights.shape[0]
+        or candidate_attribute_rewards.shape[-1] != shared_gate_weights.shape[-1]
+    ):
+        raise ValueError(
+            "Candidate rewards and shared gate weights have incompatible "
+            "pair or attribute dimensions."
+        )
+    return torch.sum(
+        candidate_attribute_rewards * shared_gate_weights.unsqueeze(1), dim=-1
+    )
+
+
+def shared_gate_checkpoint_filename(args, model_name: str, preference_name: str, reference_name: str) -> str:
+    """Build the canonical Shared-Gate V2 checkpoint filename."""
+    from attributes import attribute_selection_suffix
+
+    defaults = {
+        "learning_rate": 0.0005, "weight_decay": 0.0, "n_hidden": 1,
+        "hidden_size": 64, "dropout": 0.1, "batch_size": 2048,
+        "logit_scale": 2.0, "domain_loss_weight": 0.25,
+        "entropy_weight": 0.02, "load_balance_weight": 0.05,
+    }
+    hyperparameters = "".join(
+        f"_{key[:2]}{getattr(args, key, default)}"
+        for key, default in defaults.items()
+    )
+    debiasing_dims = (
+        [-1] if str(reference_name).lower() == "null"
+        else getattr(args, "debiasing_dims", [-1])
+    )
+    suffix = debiasing_checkpoint_suffix(
+        debiasing_dims, getattr(args, "corr_threshold", 0.04)
+    )
+    suffix += "_cv" if getattr(args, "curriculum", False) else ""
+    suffix += "_bd" if getattr(args, "balance_difficulties", False) else ""
+    suffix += "" if getattr(args, "balance_domains", True) else "_ubd"
+    suffix += "_lgs" if getattr(args, "learnable_logit_scale", False) else ""
+    entropy_floor = getattr(args, "entropy_floor_fraction", 0.35)
+    suffix += "" if entropy_floor == 0.35 else f"_ef{entropy_floor}"
+    suffix += attribute_selection_suffix(
+        getattr(args, "attribute_subset", "full"),
+        getattr(args, "exclude_attributes", []),
+    )
+    gate_input_mode = getattr(args, "gate_input_mode", "prompt")
+    gate_mode_codes = {
+        "prompt": "prompt",
+        "global": "global",
+        "shuffled_prompt": "shuffle",
+        "candidate_conditioned": "candidate",
+    }
+    if gate_input_mode not in gate_mode_codes:
+        raise ValueError(f"Unknown gate_input_mode: {gate_input_mode}")
+    if gate_input_mode != "prompt":
+        suffix += f"_gim-{gate_mode_codes[gate_input_mode]}"
+    held_out_domain = getattr(args, "held_out_domain", None)
+    if held_out_domain:
+        suffix += f"_holdout-{held_out_domain}"
+    checkpoint_tag = getattr(args, "checkpoint_tag", None)
+    suffix += f"_tag-{checkpoint_tag}" if checkpoint_tag else ""
+    validation_manifest = getattr(args, "validation_group_ids_path", None)
+    if validation_manifest:
+        import hashlib
+        with open(validation_manifest, "rb") as stream:
+            suffix += "_vs-" + hashlib.sha256(stream.read()).hexdigest()[:16]
+    suffix += "_refit" if getattr(args, "train_on_all", False) else ""
+    filename = (
+        f"gating_network_sgv2_{model_name}_mo_{args.multi_objective_dataset_name}_"
+        f"pref_{preference_name}_ref_{reference_name}"
+        f"_t{getattr(args, 'temperature', 2.0):.1f}"
+        f"_n{getattr(args, 'n_steps', 30000)}"
+        f"_seed{getattr(args, 'seed', 0)}{hyperparameters}{suffix}.pt"
+    )
+    # Linux filesystems normally limit a single path component to 255 bytes.
+    # Keep short legacy names unchanged, but make long ablation names portable
+    # and collision resistant for both Stage 2 saving and Stage 3 lookup.
+    max_filename_bytes = 240
+    if len(filename.encode("utf-8")) > max_filename_bytes:
+        import hashlib
+
+        digest = hashlib.sha256(filename.encode("utf-8")).hexdigest()[:16]
+        extension = ".pt"
+        budget = max_filename_bytes - len(f"_h{digest}{extension}")
+        filename = f"{filename[:-len(extension)][:budget]}_h{digest}{extension}"
+    return filename
+
+
+
 # ---------------------------------------------------------------------------
 # Tokenizer loading
 # ---------------------------------------------------------------------------
@@ -159,10 +298,6 @@ def _resolve_inference_model_path(
     if not isinstance(inference_cfg, dict):
         inference_cfg = {}
 
-    explicit_model_path = inference_cfg.get("model_path")
-    if explicit_model_path:
-        return str(explicit_model_path)
-
     if cli_model_parent_dir or cli_model_name:
         model_parent_dir = str(cli_model_parent_dir or inference_cfg.get("model_parent_dir", "model"))
         model_name = cli_model_name or inference_cfg.get("model_name")
@@ -170,6 +305,10 @@ def _resolve_inference_model_path(
             raise ValueError("model_name must be provided via --model_name or config.yaml inference.model_name")
         return os.path.join(model_parent_dir, str(model_name))
 
+    explicit_model_path = inference_cfg.get("model_path")
+    if explicit_model_path:
+        return str(explicit_model_path)
+
     model_name = inference_cfg.get("model_name")
     if not model_name:
         raise ValueError("model_name must be provided via --model_name or config.yaml inference.model_name")