converter fetch on recipe level

src/fairseq2/cli/_setup.py

@@ -62,15 +62,15 @@
     CausalLMLossEvalConfig,
     CausalLMTrainConfig,
     InstructionFinetuneConfig,
-    POFinetuneConfig,
     OnlineFinetuneConfig,
+    POFinetuneConfig,
     TextGenerateConfig,
     load_clm_loss_evaluator,
     load_clm_trainer,
     load_instruction_finetuner,
+    load_online_finetuner,
     load_po_finetuner,
     load_text_generator,
-    load_online_finetuner,
 )
 from fairseq2.recipes.mt import (
     MTEvalConfig,

src/fairseq2/data/text/tokenizers/huggingface_tokenizer.py

@@ -7,20 +7,20 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
 from pathlib import Path
 from typing import final
 
 import torch
 from torch import Tensor
+from transformers import AutoTokenizer
 from typing_extensions import override
 
 from fairseq2.data import VocabularyInfo
 from fairseq2.data.text.tokenizers import (
     TextTokenDecoder,
     TextTokenEncoder,
 )
 from fairseq2.typing import Device
-from transformers import AutoTokenizer
 
 
 @final

src/fairseq2/datasets/prompt.py

@@ -11,12 +11,12 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Final, List, cast, final
 
 import torch
 from typing_extensions import override
 
 from fairseq2.data import (
     CollateOptionsOverride,
     Collater,
     DataPipeline,
@@ -26,13 +26,13 @@
     read_sequence,
 )
 from fairseq2.data.text.tokenizers import TextTokenizer
 from fairseq2.datasets import (
     DataPipelineReader,
     DataReadOptions,
     DatasetHubAccessor,
+    DatasetLoadError,
     LengthBatching,
     StaticBatching,
-    DatasetLoadError,
     UnknownSplitError,
 )
 from fairseq2.datasets.utils._manifest import _load_files_and_weights
@@ -74,11 +74,11 @@
     def batch_size(self) -> int:
         """The size of the batch dimension."""
         return len(self.prompts)
-    
+
     @property
     def prompt_lengths(self) -> list[int]:
         return [len(p) for p in self.prompts]
-    
+
     @override
     def to(self, device: Device, *, non_blocking: bool = False) -> None:
         # no device moving since we only carry tokens prompts here

src/fairseq2/models/llama/_hg.py

@@ -9,10 +9,14 @@
 from pathlib import Path
 from typing import cast
 
-from torch import Tensor
 import torch
+from torch import Tensor
 
-from fairseq2.models.utils.checkpoint import convert_checkpoint, create_reverse_key_map, get_converted_key
+from fairseq2.models.utils.checkpoint import (
+    convert_checkpoint,
+    create_reverse_key_map,
+    get_converted_key,
+)
 from fairseq2.models.utils.hg import save_hg_checkpoint
 
 # isort: split
@@ -119,8 +123,8 @@ def _convert_to_hg_config(config: LLaMAConfig) -> dict[str, object]:
     }
 
 
-def _convert_parameter(name: str,
-    parameter: torch.nn.Parameter, config: LLaMAConfig
+def _convert_parameter(
+    name: str, parameter: torch.nn.Parameter, config: LLaMAConfig
 ) -> dict[str, object]:
     head_dim = config.model_dim // config.num_attn_heads
 
@@ -159,4 +163,4 @@ def permute_rotary(w: Tensor, num_heads: int) -> Tensor:
 
     converted_name = get_converted_key(name, key_map)
 
-    return converted_name, parameter
+    return converted_name, parameter

src/fairseq2/models/qwen/_hg.py

@@ -10,7 +10,11 @@
 
 import torch
 
-from fairseq2.models.utils.checkpoint import convert_checkpoint, create_reverse_key_map, get_converted_key
+from fairseq2.models.utils.checkpoint import (
+    convert_checkpoint,
+    create_reverse_key_map,
+    get_converted_key,
+)
 from fairseq2.models.utils.hg import save_hg_checkpoint
 
 # isort: split
@@ -61,8 +65,8 @@ def _convert_to_hg_checkpoint(
     return hg_checkpoint
 
 
-def _convert_parameter(name: str,
-    parameter: torch.nn.Parameter, config: QwenConfig
+def _convert_parameter(
+    name: str, parameter: torch.nn.Parameter, config: QwenConfig
 ) -> dict[str, object]:
 
     key_map = {
@@ -86,4 +90,4 @@ def _convert_parameter(name: str,
 
     converted_name = get_converted_key(name, key_map)
 
-    return converted_name, parameter
+    return converted_name, parameter

src/fairseq2/models/utils/checkpoint.py

@@ -94,13 +94,15 @@ def load_checkpoint(
     if errors:
         raise ValueError(" ".join(errors))
 
+
 def get_converted_key(key: str, key_map: Mapping[str, str]) -> str:
     for pattern, replacement in key_map.items():
         if (converted_key := re.sub(pattern, replacement, key)) != key:
             return converted_key
 
     return key
 
+
 def convert_checkpoint(
     checkpoint: dict[str, object], key_map: Mapping[str, str]
 ) -> dict[str, object]:

src/fairseq2/recipes/lm/__init__.py

@@ -39,6 +39,82 @@
 from fairseq2.recipes.lm._loss_eval import (
     register_clm_loss_eval_configs as register_clm_loss_eval_configs,
 )
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    GeneralVerifierExtractorHandler as GeneralVerifierExtractorHandler,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1PairwiseScoreExtractor as J1PairwiseScoreExtractor,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1PairwiseScoreExtractorHandler as J1PairwiseScoreExtractorHandler,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1PointwiseExtractor as J1PointwiseExtractor,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    J1PointwiseExtractorHandler as J1PointwiseExtractorHandler,
+)
+from fairseq2.recipes.lm._online_finetune._generative_judge import (
+    JudgmentExtractorHandler as JudgmentExtractorHandler,
+)
+
+# from fairseq2.recipes.lm._online_finetune._group_dpo import (
+#     GroupDpoFinetuneUnitHandler as GroupDpoFinetuneUnitHandler,
+# )
+from fairseq2.recipes.lm._online_finetune._grpo import (
+    GrpoFinetuneUnitHandler as GrpoFinetuneUnitHandler,
+)
+from fairseq2.recipes.lm._online_finetune._online_dpo import (
+    OnlineDpoFinetuneUnitHandler as OnlineDpoFinetuneUnitHandler,
+)
+from fairseq2.recipes.lm._online_finetune._recipe import (
+    OnlineFinetuneConfig,
+)
+from fairseq2.recipes.lm._online_finetune._recipe import (
+    OnlineFinetuneDatasetSection as OnlineFinetuneDatasetSection,
+)
+from fairseq2.recipes.lm._online_finetune._recipe import (
+    OnlineFinetuneUnitHandler,
+    load_online_finetuner,
+    register_online_finetune_configs,
+)
+from fairseq2.recipes.lm._online_finetune._remote_model import (
+    NoEnvAtheneRewardPipeline as NoEnvAtheneRewardPipeline,
+)
+from fairseq2.recipes.lm._online_finetune._remote_model import (
+    NoEnvGeneralVerifierPipeline as NoEnvGeneralVerifierPipeline,
+)
+from fairseq2.recipes.lm._online_finetune._remote_model import (
+    RemoteModelHandler as RemoteModelHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AtheneVerifier as AtheneVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    AtheneVerifierHandler as AtheneVerifierHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePairwiseVerifier as GenerativePairwiseVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePairwiseVerifierHandler as GenerativePairwiseVerifierHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePointwiseVerifier as GenerativePointwiseVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePointwiseVerifierHandler as GenerativePointwiseVerifierHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import GSM8kVerifier as GSM8kVerifier
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GSM8kVerifierHandler as GSM8kVerifierHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    MathVerifyHandler as MathVerifyHandler,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    VLLMOutputRewardHandler as VLLMOutputRewardHandler,
+)
 from fairseq2.recipes.lm._preference_finetune._config import (
     POCriterionSection as POCriterionSection,
 )
@@ -119,100 +195,10 @@
 from fairseq2.recipes.lm._text_generate import (
     register_text_generate_configs as register_text_generate_configs,
 )
-from fairseq2.recipes.lm._online_finetune._recipe import (
-    register_online_finetune_configs,
-    OnlineFinetuneUnitHandler,
-    load_online_finetuner,
-    OnlineFinetuneConfig,
-)
-
-from fairseq2.recipes.lm._online_finetune._recipe import (
-    OnlineFinetuneDatasetSection as OnlineFinetuneDatasetSection,
-)
-from fairseq2.recipes.lm._online_finetune._online_dpo import (
-    OnlineDpoFinetuneUnitHandler as OnlineDpoFinetuneUnitHandler,
-)
-
-# from fairseq2.recipes.lm._online_finetune._group_dpo import (
-#     GroupDpoFinetuneUnitHandler as GroupDpoFinetuneUnitHandler,
-# )
-from fairseq2.recipes.lm._online_finetune._grpo import (
-    GrpoFinetuneUnitHandler as GrpoFinetuneUnitHandler,
-)
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    VLLMOutputRewardHandler as VLLMOutputRewardHandler,
-)
-from fairseq2.recipes.lm._online_finetune._rewards import GSM8kVerifier as GSM8kVerifier
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    GSM8kVerifierHandler as GSM8kVerifierHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    AtheneVerifier as AtheneVerifier,
-)
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    AtheneVerifierHandler as AtheneVerifierHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    MathVerifyHandler as MathVerifyHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    GenerativePointwiseVerifier as GenerativePointwiseVerifier,
-)
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    GenerativePointwiseVerifierHandler as GenerativePointwiseVerifierHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    GenerativePairwiseVerifier as GenerativePairwiseVerifier,
-)
-from fairseq2.recipes.lm._online_finetune._rewards import (
-    GenerativePairwiseVerifierHandler as GenerativePairwiseVerifierHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._remote_model import (
-    RemoteModelHandler as RemoteModelHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._remote_model import (
-    NoEnvAtheneRewardPipeline as NoEnvAtheneRewardPipeline,
-)
-
-from fairseq2.recipes.lm._online_finetune._remote_model import (
-    NoEnvGeneralVerifierPipeline as NoEnvGeneralVerifierPipeline,
-)
-
-
 from fairseq2.recipes.lm._train import CausalLMTrainConfig as CausalLMTrainConfig
 from fairseq2.recipes.lm._train import CausalLMTrainUnit as CausalLMTrainUnit
 from fairseq2.recipes.lm._train import TextDatasetSection as TextDatasetSection
 from fairseq2.recipes.lm._train import load_clm_trainer as load_clm_trainer
 from fairseq2.recipes.lm._train import (
     register_clm_train_configs as register_clm_train_configs,
 )
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    JudgmentExtractorHandler as JudgmentExtractorHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    J1PointwiseExtractor as J1PointwiseExtractor,
-)
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    J1PointwiseExtractorHandler as J1PointwiseExtractorHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    J1PairwiseScoreExtractor as J1PairwiseScoreExtractor,
-)
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    J1PairwiseScoreExtractorHandler as J1PairwiseScoreExtractorHandler,
-)
-
-from fairseq2.recipes.lm._online_finetune._generative_judge import (
-    GeneralVerifierExtractorHandler as GeneralVerifierExtractorHandler,
-)

src/fairseq2/recipes/lm/_online_finetune/_common.py

@@ -17,24 +17,22 @@
 from torch import Tensor
 from vllm import RequestOutput
 
-
 from fairseq2.data import (
     CollateOptionsOverride,
     Collater,
     SequenceData,
 )
-from fairseq2.datasets.preference import PreferenceBatch
-from fairseq2.datasets.prompt import PromptBatch
-from fairseq2.gang import Gang, Gangs
 from fairseq2.datasets import (
     SequenceBatch,
 )
+from fairseq2.datasets.preference import PreferenceBatch
+from fairseq2.datasets.prompt import PromptBatch
+from fairseq2.gang import Gang, Gangs
+from fairseq2.logging import log
+from fairseq2.metrics import Mean, MetricBag, Sum
 from fairseq2.nn._batch_layout import BatchLayout
 from fairseq2.nn.utils.padding import pad_seqs
-
-from fairseq2.logging import log
 from fairseq2.recipes.lm._online_finetune._remote_model import RemoteVllmModel
-from fairseq2.metrics import Mean, Sum, MetricBag
 
 
 @dataclass(kw_only=True)
@@ -595,3 +593,18 @@ def compute_reference_logps(
     ).seqs
 
     return ref_logps
+
+
+def get_parameter_converter(model_config):
+
+    from fairseq2.models.llama import LLaMAConfig
+    from fairseq2.models.qwen import QwenConfig
+
+    if isinstance(model_config, QwenConfig):
+        from fairseq2.models.qwen._hg import _convert_parameter
+    elif isinstance(model_config, LLaMAConfig):
+        from fairseq2.models.llama._hg import _convert_parameter
+    else:
+        raise RuntimeError(f"{model_config} not supported in online recipe")
+
+    return _convert_parameter

src/fairseq2/recipes/lm/_online_finetune/_generative_judge.py

@@ -72,16 +72,18 @@
 """
 
 
+import re
 from abc import ABC, abstractmethod
+from typing import Any
+
 from typing_extensions import override
+
 from fairseq2.logging import log
-from typing import Any
-import re
 
 
 class JudgmentExtractorHandler(ABC):
     @abstractmethod
     def create(self): ...
 
     @property
     @abstractmethod
@@ -111,7 +113,7 @@
     def prompt(self) -> str: ...
 
     @abstractmethod
     def format_prompt(self, prompt_text, **kwargs: Any) -> str: ...
 
     """
     Format the prompt text and additional arguments into a string suitable for input to the reward model.
@@ -124,7 +126,7 @@
     """
 
     @abstractmethod
     def extract(self, generation) -> float | str: ...
 
     """
     Extract the final scalar reward score from the model's response.
@@ -143,7 +145,7 @@
     """
 
     @abstractmethod
     def aggregate(self, judgments) -> float | str: ...
 
     """
     Aggregate multiple responses (judgments) from the reward model into a single value.
@@ -159,31 +161,31 @@
 
 
 class GeneralVerifierExtractorHandler(JudgmentExtractorHandler):
     def __init__(self):
         pass
 
     @override
     def create(self):
         return GeneralVerifierExtractor()
 
     @property
     @override
     def name(self):
         return "general_verifier_extractor"
 
     @property
     @override
     def config_kls(self):
         return None
 
 
 class GeneralVerifierExtractor(JudgmentExtractor):
     def __init__(self):
         try:
             from math_verify import parse
             from math_verify.parser import (
-                LatexExtractionConfig,
                 ExprExtractionConfig,
+                LatexExtractionConfig,
                 NormalizationConfig,
             )
         except ImportError: