From 22d020320f97886445ac30e3475cb1d01a66c802 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Sun, 15 Mar 2026 14:21:51 -0700 Subject: [PATCH 1/2] Fix poem finetuning example and Tinker compatibility --- examples/poem_finetuning_example.py | 35 ++++++++---- .../finetune/trainer/tinker_trainer.py | 46 ++++++++++++--- tests/finetune/test_trainer.py | 56 +++++++++++++++++++ 3 files changed, 116 insertions(+), 21 deletions(-) diff --git a/examples/poem_finetuning_example.py b/examples/poem_finetuning_example.py index aa14c3c3..830093e6 100644 --- a/examples/poem_finetuning_example.py +++ b/examples/poem_finetuning_example.py @@ -6,9 +6,12 @@ 3. Fine-tuning a model using TinkerTrainer Usage: + # Install the optional fine-tuning dependency for real Tinker training + poetry install --extras finetune + # Set your API keys export OPENAI_API_KEY="your-openai-key" - export TINKER_API_KEY="your-tinker-key" # Optional, runs in mock mode without it + export TINKER_API_KEY="your-tinker-key" # Optional, runs in mock mode without it or without the SDK # Run the example poetry run python examples/poem_finetuning_example.py @@ -18,9 +21,17 @@ """ import argparse +import sys +from pathlib import Path from pydantic import BaseModel, Field +# Avoid shadowing the real `tinker` package with `examples/tinker/` when this +# file is executed as `python examples/poem_finetuning_example.py`. +EXAMPLES_DIR = Path(__file__).resolve().parent +if str(EXAMPLES_DIR) in sys.path: + sys.path.remove(str(EXAMPLES_DIR)) + from bespokelabs.curator import LLM, TinkerTrainer, TinkerTrainerConfig @@ -46,10 +57,12 @@ class PoemRequest(BaseModel): # ============================================================================= # Step 2: Create the Curator LLM for data generation # ============================================================================= -def create_poem_curator() -> LLM: - """Create a Curator LLM configured for poem generation.""" +class PoemCurator(LLM): + """Curator LLM configured for poem generation.""" + + response_format = Poem - def prompt_func(row: dict) -> str: + def prompt(self, row: dict) -> str: """Generate a prompt for creating a poem.""" return f"""You are a creative poet. Write a {row['style']} poem about "{row['theme']}". @@ -58,7 +71,7 @@ def prompt_func(row: dict) -> str: After the poem, provide a brief explanation of its meaning and imagery.""" - def parse_func(row: dict, response: Poem) -> dict: + def parse(self, row: dict, response: Poem) -> dict: """Parse the response into training format.""" return { "theme": response.theme, @@ -69,12 +82,10 @@ def parse_func(row: dict, response: Poem) -> dict: "request_style": row["style"], } - return LLM( - model_name="gpt-4o-mini", - prompt_func=prompt_func, - parse_func=parse_func, - response_format=Poem, - ) + +def create_poem_curator() -> LLM: + """Create a Curator LLM configured for poem generation.""" + return PoemCurator(model_name="gpt-4o-mini") # ============================================================================= @@ -119,7 +130,7 @@ def generate_poem_dataset() -> list: poem_response = poem_curator(requests) # Convert response to list - poem_data = poem_response.to_list() + poem_data = poem_response.dataset.to_list() print(f"Generated {len(poem_data)} poems") return poem_data diff --git a/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py b/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py index 9accbec4..b7a8fe69 100644 --- a/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py +++ b/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py @@ -1,5 +1,6 @@ """TinkerTrainer implementation for LoRA fine-tuning via Tinker API.""" +import inspect import time from typing import Any, Dict, List, Optional @@ -88,16 +89,26 @@ def _initialize_client(self) -> None: try: self._service_client = tinker.ServiceClient(api_key=self.config.api_key) - self._training_client = self._service_client.create_lora_training_client( - base_model=self.config.base_model, - rank=self.config.lora_config.rank, - alpha=self.config.lora_config.alpha, - dropout=self.config.lora_config.dropout, - target_modules=self.config.lora_config.target_modules, - train_mlp=True, - train_attn=True, - train_unembed=True, + training_client_kwargs = { + "base_model": self.config.base_model, + "rank": self.config.lora_config.rank, + "alpha": self.config.lora_config.alpha, + "dropout": self.config.lora_config.dropout, + "target_modules": self.config.lora_config.target_modules, + "train_mlp": True, + "train_attn": True, + "train_unembed": True, + } + supported_kwargs, ignored_kwargs = self._get_supported_kwargs( + self._service_client.create_lora_training_client, + training_client_kwargs, ) + if ignored_kwargs: + logger.warning( + "Installed Tinker SDK does not support LoRA args %s; continuing with supported args only.", + ", ".join(ignored_kwargs), + ) + self._training_client = self._service_client.create_lora_training_client(**supported_kwargs) self._tokenizer = self._training_client.get_tokenizer() logger.info(f"TinkerTrainer initialized for model: {self.config.base_model}") @@ -109,6 +120,18 @@ def _initialize_client(self) -> None: self._training_client = None self._tokenizer = None + @staticmethod + def _get_supported_kwargs(callable_obj: Any, kwargs: Dict[str, Any]) -> tuple[Dict[str, Any], List[str]]: + """Filter keyword arguments to those accepted by the callable.""" + try: + parameters = inspect.signature(callable_obj).parameters + except (TypeError, ValueError): + return kwargs, [] + + supported_kwargs = {key: value for key, value in kwargs.items() if key in parameters} + ignored_kwargs = [key for key in kwargs if key not in supported_kwargs] + return supported_kwargs, ignored_kwargs + def save_checkpoint(self, name: str, step: int, epoch: int, loss: float) -> Optional[CheckpointInfo]: """Save a training checkpoint. @@ -488,6 +511,11 @@ def _get_effective_target_weight(batch_data: List[Any]) -> float: if weights is None: continue + if hasattr(weights, "data"): + weights = weights.data + elif isinstance(weights, dict) and "data" in weights: + weights = weights["data"] + found_weights = True effective_weight += sum(float(weight) for weight in weights) diff --git a/tests/finetune/test_trainer.py b/tests/finetune/test_trainer.py index d850397e..b2f05a9d 100644 --- a/tests/finetune/test_trainer.py +++ b/tests/finetune/test_trainer.py @@ -151,6 +151,43 @@ def test_get_sampling_client(self, trainer): # In mock mode, client is None assert client is None + def test_initialize_client_filters_unsupported_lora_args(self, monkeypatch): + """Test client initialization against SDKs with narrower LoRA signatures.""" + captured = {} + + class FakeTrainingClient: + def get_tokenizer(self): + return SimpleNamespace(vocab_size=123) + + class FakeServiceClient: + def __init__(self, api_key): + captured["api_key"] = api_key + + def create_lora_training_client(self, base_model, rank, train_mlp=True, train_attn=True, train_unembed=True): + captured["kwargs"] = { + "base_model": base_model, + "rank": rank, + "train_mlp": train_mlp, + "train_attn": train_attn, + "train_unembed": train_unembed, + } + return FakeTrainingClient() + + monkeypatch.setattr(tinker_trainer_module, "TINKER_AVAILABLE", True) + monkeypatch.setattr(tinker_trainer_module, "tinker", SimpleNamespace(ServiceClient=FakeServiceClient)) + + trainer = TinkerTrainer(TinkerTrainerConfig(base_model="Qwen3-8B", epochs=1, api_key="test-key")) + + assert captured["api_key"] == "test-key" + assert captured["kwargs"] == { + "base_model": "Qwen3-8B", + "rank": 16, + "train_mlp": True, + "train_attn": True, + "train_unembed": True, + } + assert trainer._tokenizer.vocab_size == 123 + def test_training_step_normalizes_loss_by_weighted_targets(self, trainer, monkeypatch): """Test loss uses active target weights, not total prompt length.""" @@ -185,6 +222,25 @@ def __init__(self, input_tokens, weights): assert loss == pytest.approx(1.0) assert tokens_processed == 4 + def test_effective_target_weight_handles_tensor_wrapped_weights(self): + """Test weight extraction from tensor-like wrappers used by Tinker.""" + + class FakeTensorData: + def __init__(self, data): + self.data = data + + def __iter__(self): + return iter([("data", self.data)]) + + datum = SimpleNamespace( + model_input=SimpleNamespace(chunks=[SimpleNamespace(tokens=[10, 11, 12])]), + loss_fn_inputs={"weights": FakeTensorData([0.0, 0.5, 1.0])}, + ) + + effective_weight = TinkerTrainer._get_effective_target_weight([datum]) + + assert effective_weight == pytest.approx(1.5) + class TestCustomTrainer: """Tests for custom trainer implementations.""" From 692d390746ae18ac6f3a5c6f917bf9ff250db719 Mon Sep 17 00:00:00 2001 From: Maheswaran Sathiamoorthy Date: Sun, 15 Mar 2026 14:45:37 -0700 Subject: [PATCH 2/2] Fix Tinker kwargs handling and poem example lint --- examples/poem_finetuning_example.py | 8 ----- .../finetune/trainer/tinker_trainer.py | 15 ++++++-- tests/finetune/test_trainer.py | 34 +++++++++++++++++++ 3 files changed, 46 insertions(+), 11 deletions(-) diff --git a/examples/poem_finetuning_example.py b/examples/poem_finetuning_example.py index 830093e6..0065b987 100644 --- a/examples/poem_finetuning_example.py +++ b/examples/poem_finetuning_example.py @@ -21,17 +21,9 @@ """ import argparse -import sys -from pathlib import Path from pydantic import BaseModel, Field -# Avoid shadowing the real `tinker` package with `examples/tinker/` when this -# file is executed as `python examples/poem_finetuning_example.py`. -EXAMPLES_DIR = Path(__file__).resolve().parent -if str(EXAMPLES_DIR) in sys.path: - sys.path.remove(str(EXAMPLES_DIR)) - from bespokelabs.curator import LLM, TinkerTrainer, TinkerTrainerConfig diff --git a/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py b/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py index b7a8fe69..8c95ce4a 100644 --- a/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py +++ b/src/bespokelabs/curator/finetune/trainer/tinker_trainer.py @@ -18,10 +18,16 @@ ) from bespokelabs.curator.log import logger +tinker = None try: - import tinker - - TINKER_AVAILABLE = True + import tinker as _tinker + + required_tinker_attrs = ("ServiceClient", "Datum", "ModelInput", "EncodedTextChunk") + if all(hasattr(_tinker, attr) for attr in required_tinker_attrs): + tinker = _tinker + TINKER_AVAILABLE = True + else: + TINKER_AVAILABLE = False except ImportError: TINKER_AVAILABLE = False @@ -128,6 +134,9 @@ def _get_supported_kwargs(callable_obj: Any, kwargs: Dict[str, Any]) -> tuple[Di except (TypeError, ValueError): return kwargs, [] + if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()): + return kwargs, [] + supported_kwargs = {key: value for key, value in kwargs.items() if key in parameters} ignored_kwargs = [key for key in kwargs if key not in supported_kwargs] return supported_kwargs, ignored_kwargs diff --git a/tests/finetune/test_trainer.py b/tests/finetune/test_trainer.py index b2f05a9d..9dd24374 100644 --- a/tests/finetune/test_trainer.py +++ b/tests/finetune/test_trainer.py @@ -188,6 +188,40 @@ def create_lora_training_client(self, base_model, rank, train_mlp=True, train_at } assert trainer._tokenizer.vocab_size == 123 + def test_initialize_client_keeps_kwargs_for_var_keyword_signatures(self, monkeypatch): + """Test client initialization preserves kwargs when the SDK accepts **kwargs.""" + captured = {} + + class FakeTrainingClient: + def get_tokenizer(self): + return SimpleNamespace(vocab_size=123) + + class FakeServiceClient: + def __init__(self, api_key): + captured["api_key"] = api_key + + def create_lora_training_client(self, **kwargs): + captured["kwargs"] = kwargs + return FakeTrainingClient() + + monkeypatch.setattr(tinker_trainer_module, "TINKER_AVAILABLE", True) + monkeypatch.setattr(tinker_trainer_module, "tinker", SimpleNamespace(ServiceClient=FakeServiceClient)) + + trainer = TinkerTrainer(TinkerTrainerConfig(base_model="Qwen3-8B", epochs=1, api_key="test-key")) + + assert captured["api_key"] == "test-key" + assert captured["kwargs"] == { + "base_model": "Qwen3-8B", + "rank": 16, + "alpha": 32, + "dropout": 0.05, + "target_modules": ["q_proj", "v_proj", "k_proj", "o_proj"], + "train_mlp": True, + "train_attn": True, + "train_unembed": True, + } + assert trainer._tokenizer.vocab_size == 123 + def test_training_step_normalizes_loss_by_weighted_targets(self, trainer, monkeypatch): """Test loss uses active target weights, not total prompt length."""