Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171"""Fine-tune a LoRA on Tinker from `machine export` JSONL.
Run inside a tinker-cookbook checkout's environment:
cd ~/code/tinker-cookbook && source .venv/bin/activate python ~/code/machine/scripts/train_sft.py data.jsonl ~/code/machine/.machine/training/my-run \\ --model Qwen/Qwen3.6-35B-A3B --epochs 2
Logging under a workspace's .machine/training/ puts the run on`machine dashboard`, which charts its losses while it trains.
A line with "whole_run": true is one whole run, trained on every assistantmessage marked "trainable" (Inkling only: each step's prompt is a token prefixof the run's sequence there, so it matches sampling at a fraction of thetokens of one example per step).
Adapted from tinker-cookbook's recipes/sl_basic.py. The learning rate defaultsto the cookbook's calibrated rate for the model (Llama and Qwen); Inkling hasnone, so it defaults to 2e-4, the rate of the cookbook's Inkling SFT recipe(recipes/audio/asr/sl_train.py). Inkling (the tml_v0 renderer) is rendered at--effort, which must be the thinking effort runs sampled at: Tinker's OpenAIendpoint defaults to high, 0.9. Trainson the last assistant message of each conversation, which is correct for theQwen3.5/3.6 renderer on multi-turn data and identical on single-turn data.Checkpoints, including the sampler path to evaluate with`machine suite --model tinker://...`, are written to <log_path>/checkpoints.jsonl,and per-step losses to <log_path>/metrics.jsonl."""
import argparseimport asyncioimport jsonimport osimport shutil
import chzfrom tinker_cookbook import hyperparam_utils, model_infofrom tinker_cookbook.renderers import TrainOnWhatfrom tinker_cookbook.renderers.base import ToolCallfrom tinker_cookbook.supervised.common import datum_from_model_input_weightsfrom tinker_cookbook.supervised.data import conversation_to_datumfrom tinker_cookbook.supervised import trainfrom tinker_cookbook.supervised.data import FromConversationFileBuilderfrom tinker_cookbook.supervised.types import ChatDatasetBuilderCommonConfig
parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])parser.add_argument("data", help="JSONL from `machine export`")parser.add_argument("log_path", help="directory for checkpoints and metrics, e.g. .machine/training/<name>")parser.add_argument("--model", default="Qwen/Qwen3.6-35B-A3B")parser.add_argument("--epochs", type=int, default=2)parser.add_argument("--batch-size", type=int, default=8)parser.add_argument("--lr", type=float, default=None, help="default: the cookbook's calibrated rate")parser.add_argument("--test-size", type=int, default=8, help="examples held out to track loss")parser.add_argument("--max-length", type=int, default=65536)parser.add_argument("--dry-run", action="store_true", help="render the first batch, print its token counts, and exit without training")parser.add_argument("--effort", type=float, default=0.9, help="Inkling only: thinking effort to render at; match the runs' (endpoint default: high, 0.9)")args = parser.parse_args()
# Keep the training data with the run, so the run can be inspected and reproduced later.os.makedirs(args.log_path, exist_ok=True)data_copy = os.path.join(args.log_path, "data.jsonl")if os.path.abspath(args.data) != os.path.abspath(data_copy): shutil.copyfile(args.data, data_copy)
renderer = model_info.get_recommended_renderer_name(args.model)if args.lr is None: try: args.lr = hyperparam_utils.get_lr(args.model) except NotImplementedError: if not args.model.startswith("thinkingmachines/Inkling"): raise args.lr = 2e-4 print(f"learning rate {args.lr}")
@chz.chzclass ToolAwareConversationFileBuilder(FromConversationFileBuilder): """FromConversationFileBuilder for `machine export` lines that carry tools.
A line's "tools" (the tools the model was offered) are rendered into the system prompt with the renderer's own tool declaration, as the sampling endpoint did, and OpenAI-style tool_calls become the renderer's ToolCall objects. Lines without tools are rendered as before. """
def __call__(self): train_ds, test_ds = super().__call__() renderer_obj = self.renderer max_length = self.common_config.max_length train_on_what = TrainOnWhat(self.common_config.train_on_what)
def convert(row: dict) -> list: messages = [dict(m) for m in row["messages"]] for m in messages: for key in [k for k, v in m.items() if v is None]: del m[key] if m.get("tool_calls"): m["tool_calls"] = [ ToolCall(id=tc["id"], function=ToolCall.FunctionBody(name=tc["function"]["name"], arguments=tc["function"]["arguments"])) for tc in m["tool_calls"] ] tools = json.loads(row["tools_json"]) if row.get("tools_json") else None if tools and messages and messages[0]["role"] == "system": system = messages[0]["content"] text = system if isinstance(system, str) else "".join(p.get("text", "") for p in system if p.get("type") == "text") # Qwen's renderer declares the tools as given, so they go in wrapped as OpenAI tools # ({"type": "function", "function": ...}), which is how Tinker's OpenAI-compatible endpoint # renders them. tml_v0 (Inkling) wraps them itself. Either way the prompt is token for token # what the endpoint sampled from (checked against its prompt_tokens on dev runs). declared = tools if renderer == "tml_v0" else [{"type": "function", "function": t} for t in tools] messages = renderer_obj.create_conversation_prefix_with_tools(declared, system_prompt=text) + messages[1:] return messages
def map_fn(row: dict): if row.get("whole_run") and renderer != "tml_v0": raise ValueError(f"whole-run examples are checked for Inkling's tml_v0 renderer only, not {renderer}; train on step examples") if renderer == "tml_v0": # A whole-run example trains every assistant message marked trainable; a step example, its last. what = TrainOnWhat.CUSTOMIZED if row.get("whole_run") else train_on_what model_input, weights = renderer_obj.build_supervised_example(convert(row), train_on_what=what, effort=args.effort) return datum_from_model_input_weights(model_input, weights, max_length, reduction="mean") return conversation_to_datum(convert(row), renderer_obj, max_length, train_on_what)
for ds in (train_ds, test_ds): if ds is not None: ds.map_fn = map_fn return train_ds, test_ds
# Tools go in as a JSON string: their nested schemas do not fit one Arrow column type across lines.data_path = args.datawith open(args.data) as f: lines = [json.loads(line) for line in f if line.strip()]if any("tools" in line for line in lines): data_path = os.path.join(args.log_path, "data.tools.jsonl") with open(data_path, "w") as f: for line in lines: tools = line.pop("tools", None) line["tools_json"] = json.dumps(tools) if tools else "" f.write(json.dumps(line) + "\n")
dataset = ToolAwareConversationFileBuilder( common_config=ChatDatasetBuilderCommonConfig( model_name_for_tokenizer=args.model, renderer_name=renderer, max_length=args.max_length, batch_size=args.batch_size, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE, ), file_path=data_path, test_size=args.test_size,)if args.dry_run: train_ds, _ = dataset() for datum in train_ds.get_batch(0): weights = datum.loss_fn_inputs["weights"].to_numpy() print(f"example: {datum.model_input.length} tokens, {int((weights > 0).sum())} trained") raise SystemExit(0)
config = chz.Blueprint(train.Config).apply({ "log_path": args.log_path, "recipe_name": "machine_sft", "model_name": args.model, "renderer_name": renderer, "dataset_builder": dataset, "learning_rate": args.lr, "num_epochs": args.epochs, "eval_every": 5,}).make()asyncio.run(train.main(config))