diff --git a/.github/workflows/training-ci.yml b/.github/workflows/training-ci.yml new file mode 100644 index 00000000..6f6e2e18 --- /dev/null +++ b/.github/workflows/training-ci.yml @@ -0,0 +1,96 @@ +name: Training CI + +on: + pull_request: + paths: + - "training/**" + - ".github/workflows/training-ci.yml" + push: + branches: + - main + paths: + - "training/**" + - ".github/workflows/training-ci.yml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: training-ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + unit: + name: Unit And Import Tests + runs-on: ubuntu-latest + timeout-minutes: 20 + defaults: + run: + working-directory: training + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set Up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: "pip" + cache-dependency-path: training/pyproject.toml + + - name: Install Dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e '.[dev]' + + - name: Run Unit Tests + run: | + pytest -q tests/unit tests/test_smoke_imports.py examples/frozen_lake/test_masking.py + + qwen3-4b-smoke: + name: Qwen3-4B Smoke (${{ matrix.test_target }}) + needs: unit + runs-on: ubuntu-latest + timeout-minutes: 180 + strategy: + fail-fast: false + max-parallel: 1 + matrix: + test_target: + - tests/smoke_test/test_sft_smoke.py + - tests/smoke_test/test_dpo_smoke.py + - tests/smoke_test/test_grpo_smoke.py + defaults: + run: + working-directory: training + env: + FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }} + FIREWORKS_ACCOUNT_ID: ${{ secrets.FIREWORKS_ACCOUNT_ID }} + FIREWORKS_BASE_URL: ${{ vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }} + FIREWORKS_INFERENCE_URL: ${{ vars.FIREWORKS_INFERENCE_URL || vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }} + FIREWORKS_HOTLOAD_API_URL: ${{ vars.FIREWORKS_HOTLOAD_API_URL || vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }} + FIREWORKS_GATEWAY_SECRET: ${{ secrets.FIREWORKS_GATEWAY_SECRET }} + FIREWORKS_CUSTOM_IMAGE_TAG: ${{ vars.FIREWORKS_CUSTOM_IMAGE_TAG }} + FIREWORKS_SMOKE_BASE_MODEL: ${{ vars.FIREWORKS_SMOKE_BASE_MODEL || 'accounts/fireworks/models/qwen3-4b' }} + FIREWORKS_SMOKE_TOKENIZER_MODEL: ${{ vars.FIREWORKS_SMOKE_TOKENIZER_MODEL || 'Qwen/Qwen3-4B' }} + FIREWORKS_SMOKE_TRAINING_SHAPE: ${{ vars.FIREWORKS_SMOKE_TRAINING_SHAPE || 'ts-qwen3-4b-smoke-v1' }} + steps: + - name: Checkout + uses: actions/checkout@v4 + + - name: Set Up Python + uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: "pip" + cache-dependency-path: training/pyproject.toml + + - name: Install Dependencies + run: | + python -m pip install --upgrade pip + python -m pip install -e '.[dev]' + + - name: Run Smoke Test + run: | + pytest -q -s ${{ matrix.test_target }} diff --git a/training/examples/frozen_lake/assets/img/cracked_hole.png b/training/examples/frozen_lake/assets/img/cracked_hole.png new file mode 100644 index 00000000..55930420 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/cracked_hole.png differ diff --git a/training/examples/frozen_lake/assets/img/elf_down.png b/training/examples/frozen_lake/assets/img/elf_down.png new file mode 100644 index 00000000..afa3daf1 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/elf_down.png differ diff --git a/training/examples/frozen_lake/assets/img/elf_left.png b/training/examples/frozen_lake/assets/img/elf_left.png new file mode 100644 index 00000000..bc9e22ea Binary files /dev/null and b/training/examples/frozen_lake/assets/img/elf_left.png differ diff --git a/training/examples/frozen_lake/assets/img/elf_right.png b/training/examples/frozen_lake/assets/img/elf_right.png new file mode 100644 index 00000000..83640315 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/elf_right.png differ diff --git a/training/examples/frozen_lake/assets/img/elf_up.png b/training/examples/frozen_lake/assets/img/elf_up.png new file mode 100644 index 00000000..933f1f00 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/elf_up.png differ diff --git a/training/examples/frozen_lake/assets/img/goal.png b/training/examples/frozen_lake/assets/img/goal.png new file mode 100644 index 00000000..709e4b48 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/goal.png differ diff --git a/training/examples/frozen_lake/assets/img/hole.png b/training/examples/frozen_lake/assets/img/hole.png new file mode 100644 index 00000000..c7afd1db Binary files /dev/null and b/training/examples/frozen_lake/assets/img/hole.png differ diff --git a/training/examples/frozen_lake/assets/img/ice.png b/training/examples/frozen_lake/assets/img/ice.png new file mode 100644 index 00000000..95dbc74a Binary files /dev/null and b/training/examples/frozen_lake/assets/img/ice.png differ diff --git a/training/examples/frozen_lake/assets/img/stool.png b/training/examples/frozen_lake/assets/img/stool.png new file mode 100644 index 00000000..f304d797 Binary files /dev/null and b/training/examples/frozen_lake/assets/img/stool.png differ diff --git a/training/examples/frozen_lake/frozen_lake_env.py b/training/examples/frozen_lake/frozen_lake_env.py index c8f1d6bd..e9ebfbc8 100644 --- a/training/examples/frozen_lake/frozen_lake_env.py +++ b/training/examples/frozen_lake/frozen_lake_env.py @@ -8,6 +8,7 @@ from typing import Any, Dict, Iterable, List, Sequence, Tuple from training.examples.frozen_lake.frozen_lake_schema import FROZEN_LAKE_ACTIONS +from training.examples.frozen_lake.rendering import render_frozen_lake_png_data_url ActionToDelta = { "LEFT": (0, -1), @@ -150,12 +151,14 @@ def __init__(self, map_rows: Sequence[str], max_steps: int = 30): self._step_count = 0 self._terminated = False self._truncated = False + self._last_action: str | None = None def reset(self) -> Dict[str, Any]: self._position = 0 self._step_count = 0 self._terminated = False self._truncated = False + self._last_action = None return self._current_state(action="RESET", reward=0.0) def step(self, action: str) -> Dict[str, Any]: @@ -164,6 +167,7 @@ def step(self, action: str) -> Dict[str, Any]: raise ValueError(f"Invalid action '{action}'. Expected one of {FROZEN_LAKE_ACTIONS}") if self._terminated or self._truncated: + self._last_action = normalized_action return self._current_state(action=normalized_action, reward=0.0) row = self._position // self._side @@ -174,6 +178,7 @@ def step(self, action: str) -> Dict[str, Any]: self._position = row * self._side + col self._step_count += 1 + self._last_action = normalized_action tile = self._map_rows[row][col] reward = 1.0 if tile == "G" else 0.0 @@ -199,6 +204,17 @@ def _current_state(self, action: str, reward: float) -> Dict[str, Any]: ) return result.as_tool_result() + def render_image_data_url(self, *, cell_size: int = 96) -> str: + row = self._position // self._side + col = self._position % self._side + return render_frozen_lake_png_data_url( + self._map_rows, + agent_row=row, + agent_col=col, + last_action=self._last_action, + cell_size=cell_size, + ) + def build_frozen_lake_tool_env(environment_context: Dict[str, Any] | None, max_steps: int) -> FrozenLakeToolEnv: """Create a FrozenLakeToolEnv from dataset ``environment_context``.""" diff --git a/training/examples/frozen_lake/frozen_lake_rollout.py b/training/examples/frozen_lake/frozen_lake_rollout.py index 3cea16ba..6798b6eb 100644 --- a/training/examples/frozen_lake/frozen_lake_rollout.py +++ b/training/examples/frozen_lake/frozen_lake_rollout.py @@ -14,11 +14,13 @@ from types import SimpleNamespace from typing import Any, Dict, Iterable, List, Optional, Tuple +from fireworks import AsyncFireworks from openai.types import CompletionUsage from eval_protocol.integrations.fireworks_v1_completions_client import ( FireworksV1CompletionsClient, ParsedToolCall, + _normalize_token_id_sequence, strip_chat_special_tokens, to_openai_tool_calls, ) @@ -41,6 +43,538 @@ ) logger = logging.getLogger(__name__) +_TOKENIZER_CACHE: Dict[str, Any] = {} + +DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS = ( + "You are an RL policy for FrozenLake.\n" + "Pick the action that moves toward G while avoiding H.\n" + "Always respond with exactly one tool call and no text." +) + +DEFAULT_VISUAL_PROMPT_TEMPLATE = ( + "You are playing FrozenLake. The image shows the current grid. " + "The current textual observation is below.\n" + "{observation}\n\n" + "Tiles are labeled S, F, H, and G. The agent is marked with a red dot. " + "Use exactly one lake_move tool call with action LEFT, DOWN, RIGHT, or UP." +) + + +def _get_hf_tokenizer(tokenizer_name_or_path: str): + tok = _TOKENIZER_CACHE.get(tokenizer_name_or_path) + if tok is not None: + return tok + from transformers import AutoTokenizer + + tok = AutoTokenizer.from_pretrained(tokenizer_name_or_path, trust_remote_code=True) + _TOKENIZER_CACHE[tokenizer_name_or_path] = tok + return tok + + +def _is_kimi_tokenizer_name(tokenizer_name_or_path: str) -> bool: + lowered = str(tokenizer_name_or_path or "").lower() + return "kimi-k2.5" in lowered or "kimi-k2p5" in lowered + + +def _build_multimodal_image_content(*, image_url: str, text: Optional[str] = None) -> List[Dict[str, Any]]: + parts: List[Dict[str, Any]] = [ + {"type": "image_url", "image_url": {"url": str(image_url)}}, + ] + if text is not None: + parts.append({"type": "text", "text": str(text)}) + return parts + + +def _normalize_multimodal_content(content: Any) -> Any: + if content is None or isinstance(content, str): + return content + if not isinstance(content, list): + return str(content) + + normalized_parts: List[Dict[str, Any]] = [] + for part in content: + if isinstance(part, dict): + part_type = str(part.get("type") or "").strip().lower() + if part_type == "image": + normalized_parts.append({"type": "image"}) + continue + if part_type == "image_url": + image_url = part.get("image_url") + if isinstance(image_url, dict): + normalized_parts.append({"type": "image_url", "image_url": dict(image_url)}) + elif image_url is not None: + normalized_parts.append({"type": "image_url", "image_url": {"url": str(image_url)}}) + else: + normalized_parts.append({"type": "image"}) + continue + if part_type == "text": + normalized_parts.append({"type": "text", "text": str(part.get("text") or "")}) + continue + if "text" in part: + normalized_parts.append({"type": "text", "text": str(part.get("text") or "")}) + continue + normalized_parts.append({"type": "text", "text": str(part)}) + return normalized_parts + + +def _sanitize_messages_for_multimodal_template(messages: Iterable[Dict[str, Any]]) -> List[Dict[str, Any]]: + sanitized: List[Dict[str, Any]] = [] + for msg in messages: + sanitized_msg: Dict[str, Any] = { + "role": str(msg.get("role", "user")), + "content": _normalize_multimodal_content(msg.get("content")), + } + if msg.get("tool_calls") is not None: + sanitized_msg["tool_calls"] = msg.get("tool_calls") + if msg.get("tool_call_id") is not None: + sanitized_msg["tool_call_id"] = msg.get("tool_call_id") + if msg.get("name") is not None: + sanitized_msg["name"] = msg.get("name") + sanitized.append(sanitized_msg) + return sanitized + + +def _build_multimodal_fallback_prompt_text( + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], +) -> str: + chunks: List[str] = [] + if tools: + chunks.append("TOOLS:") + for tool in tools: + function = tool.get("function", {}) + chunks.append( + json.dumps( + { + "name": function.get("name"), + "description": function.get("description"), + "parameters": function.get("parameters"), + }, + ensure_ascii=False, + separators=(",", ":"), + ) + ) + chunks.append("") + + for msg in messages: + role = str(msg.get("role", "user")).upper() + content = msg.get("content") + if isinstance(content, list): + rendered_parts: List[str] = [] + for part in content: + if not isinstance(part, dict): + rendered_parts.append(str(part)) + continue + part_type = str(part.get("type") or "").strip().lower() + if part_type in {"image", "image_url"}: + rendered_parts.append("") + else: + rendered_parts.append(str(part.get("text") or "")) + rendered = "".join(rendered_parts) + else: + rendered = str(content or "") + chunks.append(f"{role}: {rendered}") + if msg.get("tool_calls"): + chunks.append(f"{role}_TOOL_CALLS: {json.dumps(msg['tool_calls'], ensure_ascii=False)}") + chunks.append("ASSISTANT:") + return "\n".join(chunks) + + +def _build_visual_user_prompt( + *, + prompt_template: Optional[str], + observation: str, + default_prompt: str, +) -> str: + template = str(prompt_template or default_prompt) + if "{observation}" in template: + return template.replace("{observation}", observation) + observation_suffix = f"\n\nCurrent textual observation:\n{observation}" if observation else "" + return f"{template.rstrip()}{observation_suffix}" + + +def _strip_trailing_text_suffix(text: str, suffix: str) -> str: + if suffix and text.endswith(suffix): + return text[: -len(suffix)] + return text + + +def _strip_trailing_token_suffix(token_ids: List[int], suffix_ids: List[int]) -> Tuple[List[int], int]: + normalized_token_ids = [int(x) for x in list(token_ids)] + normalized_suffix_ids = [int(x) for x in list(suffix_ids)] + if ( + normalized_suffix_ids + and len(normalized_token_ids) >= len(normalized_suffix_ids) + and normalized_token_ids[-len(normalized_suffix_ids):] == normalized_suffix_ids + ): + return normalized_token_ids[:-len(normalized_suffix_ids)], len(normalized_suffix_ids) + return normalized_token_ids, 0 + + +def _build_kimi_toolcall_generation_prefill_token_ids(tokenizer_name_or_path: str) -> List[int]: + if not _is_kimi_tokenizer_name(tokenizer_name_or_path): + return [] + tokenizer = _get_hf_tokenizer(tokenizer_name_or_path) + return _normalize_token_id_sequence( + tokenizer.encode("<|tool_calls_section_begin|>", add_special_tokens=False) + ) + + +def _build_kimi_toolcall_generation_prefill_text(tokenizer_name_or_path: str) -> str: + if not _is_kimi_tokenizer_name(tokenizer_name_or_path): + return "" + return "<|tool_calls_section_begin|>" + + +class FireworksV1ImageCompletionsClient: + """Multimodal `/v1/completions` client that mirrors local tokenized chat mode.""" + + def __init__( + self, + *, + model_id: str, + tokenizer_name_or_path: str, + api_key: Optional[str] = None, + base_url: Optional[str] = None, + temperature: float = 1.0, + max_tokens: int = 256, + request_params: Optional[Dict[str, Any]] = None, + logprobs: bool = False, + enable_thinking: Optional[bool] = None, + tool_call_parser=None, + default_tools: Optional[List[Dict[str, Any]]] = None, + ): + self.model_id = model_id + self.tokenizer_name_or_path = tokenizer_name_or_path or model_id + self._api_key = api_key + self._base_url = base_url + self.temperature = temperature + self.max_tokens = max_tokens + self.request_params = dict(request_params or {}) + self.logprobs = logprobs + self.enable_thinking = enable_thinking + self.tool_call_parser = tool_call_parser + self.default_tools = default_tools or [] + self._tokenizer = None + self._client = None + + async def close(self) -> None: + if self._client is not None: + await self._client.close() + + def _get_client(self) -> AsyncFireworks: + if self._client is None: + self._client = AsyncFireworks(api_key=self._api_key, base_url=self._base_url) + return self._client + + def _get_tokenizer(self): + if self._tokenizer is None: + self._tokenizer = _get_hf_tokenizer(self.tokenizer_name_or_path) + return self._tokenizer + + def _thinking_kwargs(self) -> Dict[str, Any]: + if self.enable_thinking is not None: + return {"thinking": self.enable_thinking} + return {} + + def _strip_generation_think_from_token_ids(self, token_ids: List[int]) -> List[int]: + if self.enable_thinking is not False or not _is_kimi_tokenizer_name(self.tokenizer_name_or_path): + return token_ids + tokenizer = self._get_tokenizer() + think_suffix_ids = _normalize_token_id_sequence( + tokenizer.encode("", add_special_tokens=False) + ) + if think_suffix_ids and token_ids[-len(think_suffix_ids):] == think_suffix_ids: + return token_ids[:-len(think_suffix_ids)] + return token_ids + + def _strip_generation_think_from_text(self, text: str) -> str: + if self.enable_thinking is not False or not _is_kimi_tokenizer_name(self.tokenizer_name_or_path): + return text + stripped = _strip_trailing_text_suffix(text, "") + stripped = _strip_trailing_text_suffix(stripped, "\n") + return stripped + + def _strip_generation_assistant_suffix_from_text(self, text: str) -> str: + return _strip_trailing_text_suffix(text, self.assistant_turn_suffix_text()) + + def _apply_chat_template( + self, + *, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + tokenize: bool, + add_generation_prompt: bool, + ) -> Any: + tokenizer = self._get_tokenizer() + sanitized_messages = _sanitize_messages_for_multimodal_template(messages=messages) + thinking_kw = self._thinking_kwargs() + try: + return tokenizer.apply_chat_template( + sanitized_messages, + tools=tools, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + **thinking_kw, + ) + except Exception as exc: + if tools: + logger.debug("Multimodal chat template with tools failed, retrying without tools: %s", exc) + return tokenizer.apply_chat_template( + sanitized_messages, + tokenize=tokenize, + add_generation_prompt=add_generation_prompt, + **thinking_kw, + ) + raise + + def build_prompt_token_ids( + self, + *, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + ) -> List[int]: + tokenizer = self._get_tokenizer() + try: + token_ids = self._apply_chat_template( + messages=messages, + tools=tools, + tokenize=True, + add_generation_prompt=True, + ) + return self._strip_generation_think_from_token_ids( + _normalize_token_id_sequence(token_ids) + ) + except Exception as exc: + logger.debug("Multimodal tokenized prompt build failed, using text fallback: %s", exc) + fallback_prompt = self.build_prompt_text(messages=messages, tools=tools) + return _normalize_token_id_sequence(tokenizer.encode(fallback_prompt, add_special_tokens=False)) + + def build_prompt_text( + self, + *, + messages: List[Dict[str, Any]], + tools: Optional[List[Dict[str, Any]]], + ) -> str: + try: + prompt_text = str( + self._apply_chat_template( + messages=messages, + tools=tools, + tokenize=False, + add_generation_prompt=True, + ) + ) + return self._strip_generation_think_from_text(prompt_text) + except Exception as exc: + logger.warning( + "Multimodal chat template build failed for %s: %s", + self.tokenizer_name_or_path, + exc, + ) + sanitized_messages = _sanitize_messages_for_multimodal_template(messages=messages) + return self._strip_generation_think_from_text( + _build_multimodal_fallback_prompt_text(sanitized_messages, tools=tools) + ) + + def build_tool_response_suffix_token_ids( + self, + *, + tool_message: Dict[str, Any], + ) -> List[int]: + tokenizer = self._get_tokenizer() + suffix_messages = [tool_message] + try: + token_ids = self._apply_chat_template( + messages=suffix_messages, + tools=None, + tokenize=True, + add_generation_prompt=True, + ) + return self._strip_generation_think_from_token_ids( + _normalize_token_id_sequence(token_ids) + ) + except Exception as exc: + logger.debug("Multimodal tool suffix build failed, using text fallback: %s", exc) + fallback_prompt = self.build_prompt_text(messages=suffix_messages, tools=None) + return _normalize_token_id_sequence(tokenizer.encode(fallback_prompt, add_special_tokens=False)) + + def build_tool_response_suffix_text( + self, + *, + tool_message: Dict[str, Any], + ) -> str: + suffix_messages = [tool_message] + return self.build_prompt_text(messages=suffix_messages, tools=None) + + def encode_assistant_turn_suffix(self) -> List[int]: + tokenizer = self._get_tokenizer() + suffix_text = "<|im_end|>" if _is_kimi_tokenizer_name(self.tokenizer_name_or_path) else "<|im_end|>\n" + return _normalize_token_id_sequence( + tokenizer.encode(suffix_text, add_special_tokens=False) + ) + + def assistant_turn_suffix_text(self) -> str: + return "<|im_end|>" if _is_kimi_tokenizer_name(self.tokenizer_name_or_path) else "<|im_end|>\n" + + def decode_token_ids(self, *, token_ids: List[int]) -> str: + if not token_ids: + return "" + tokenizer = self._get_tokenizer() + try: + return str( + tokenizer.decode( + token_ids, + skip_special_tokens=False, + clean_up_tokenization_spaces=False, + ) + ) + except TypeError: + return str(tokenizer.decode(token_ids)) + + async def create_completion_from_prompt_ids( + self, + *, + prompt_token_ids: List[int], + prompt_text: Optional[str] = None, + images: List[str], + tools: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + active_tools = tools if tools is not None else (self.default_tools or None) + normalized_prompt_token_ids = [int(x) for x in list(prompt_token_ids)] + if prompt_text is None: + prompt_text = self.decode_token_ids(token_ids=normalized_prompt_token_ids) + request_payload = { + **self.request_params, + "model": self.model_id, + "prompt": prompt_text, + "images": images, + "temperature": self.temperature, + "max_tokens": self.max_tokens, + "return_token_ids": True, + "raw_output": True, + } + if self.logprobs: + request_payload["logprobs"] = True + + max_retries = 40 + base_delay = 10.0 + for attempt in range(max_retries + 1): + try: + response = await self._get_client().completions.create(**request_payload) + break + except Exception as exc: + status = getattr(exc, "status_code", None) or getattr(exc, "status", None) + err_str = str(exc) + is_transient = ( + status in (425, 429, 502, 503, 504) + or "model_not_ready" in err_str + or "hot loading" in err_str + or "Model not found" in err_str + or "DEPLOYMENT_SCALING_UP" in err_str + ) + if not is_transient or attempt >= max_retries: + raise + delay = min(base_delay * (2 ** attempt), 60.0) + logger.info( + "Retryable multimodal error (attempt %d/%d, status=%s), retrying in %.1fs: %s", + attempt + 1, + max_retries, + status, + delay, + err_str[:200], + ) + await asyncio.sleep(delay) + + response_dict = response.model_dump() if hasattr(response, "model_dump") else dict(response) + choices = response_dict.get("choices") or [] + if not choices: + raise ValueError("Fireworks /v1/completions response did not include choices") + + choice = choices[0] + finish_reason = str(choice.get("finish_reason") or "unknown") + raw_output = choice.get("raw_output") if isinstance(choice.get("raw_output"), dict) else {} + raw_completion_token_ids = _normalize_token_id_sequence( + choice.get("token_ids") or raw_output.get("completion_token_ids") or [] + ) + choice_prompt_token_ids = _normalize_token_id_sequence( + choice.get("prompt_token_ids") or raw_output.get("prompt_token_ids") or normalized_prompt_token_ids + ) + completion_token_ids = self._strip_generation_think_from_token_ids(raw_completion_token_ids) + completion_token_ids, _ = _strip_trailing_token_suffix( + completion_token_ids, + self.encode_assistant_turn_suffix(), + ) + completion_text = self.decode_token_ids(token_ids=completion_token_ids) + if not completion_text: + completion_text = self._strip_generation_assistant_suffix_from_text( + self._strip_generation_think_from_text(str(choice.get("text") or "")) + ) + + completion_logprobs: List[float] = [] + choice_logprobs = choice.get("logprobs") + if isinstance(choice_logprobs, dict): + token_logprobs = choice_logprobs.get("token_logprobs") or [] + completion_logprobs = [float(lp) if lp is not None else 0.0 for lp in token_logprobs] + completion_logprobs = completion_logprobs[: len(completion_token_ids)] + + if self.tool_call_parser is not None: + parsed_output = self.tool_call_parser(completion_text, completion_token_ids, active_tools) + parsed_tool_call: Optional[ParsedToolCall] = parsed_output.get("parsed_tool_call") + assistant_content = str(parsed_output.get("assistant_content", "") or "") + parser_name = str(parsed_output.get("parser", "external")) + message_payload: Dict[str, Any] = { + "role": "assistant", + "content": assistant_content, + } + if parsed_tool_call is not None: + message_payload["tool_calls"] = to_openai_tool_calls(parsed_tool_call) + else: + assistant_content = strip_chat_special_tokens(completion_text) + parser_name = "none" + message_payload = {"role": "assistant", "content": assistant_content} + + usage_obj = response_dict.get("usage") or {} + result: Dict[str, Any] = { + "choices": [ + { + "message": message_payload, + "finish_reason": finish_reason, + "raw_output": {**dict(raw_output or {}), "tool_call_parser": parser_name}, + } + ], + "usage": { + "prompt_tokens": int(usage_obj.get("prompt_tokens", len(choice_prompt_token_ids))), + "completion_tokens": int(usage_obj.get("completion_tokens", len(completion_token_ids))), + "total_tokens": int( + usage_obj.get("total_tokens", len(choice_prompt_token_ids) + len(completion_token_ids)) + ), + }, + "prompt_text": prompt_text, + "prompt_ids": list(choice_prompt_token_ids), + "completion_ids": list(completion_token_ids), + "completion_text": completion_text, + "finish_reason": finish_reason, + "raw_output": {**dict(raw_output or {}), "tool_call_parser": parser_name}, + } + if completion_logprobs: + result["completion_logprobs"] = completion_logprobs + return result + + async def create_completion( + self, + *, + messages: List[Dict[str, Any]], + images: List[str], + tools: Optional[List[Dict[str, Any]]] = None, + ) -> Dict[str, Any]: + active_tools = tools if tools is not None else (self.default_tools or None) + prompt_token_ids = self.build_prompt_token_ids(messages=messages, tools=active_tools) + return await self.create_completion_from_prompt_ids( + prompt_token_ids=prompt_token_ids, + images=images, + tools=active_tools, + ) def _to_message_payload(messages: Iterable[Message]) -> List[Dict[str, Any]]: @@ -90,7 +624,8 @@ def _merge_request_params( ignored_keys = { "model", "temperature", "max_tokens", "base_url", "api_key", "tokenizer_name_or_path", "tokenizer_path", - "system_prompt", "user_prompt_template", + "system_prompt", "user_prompt_template", "visual_prompt_template", + "observation_mode", "allow_plaintext_action_fallback", } for key, value in completion_params.items(): @@ -103,6 +638,17 @@ def _merge_request_params( return merged +def _looks_like_tool_parse_error_message(message: str) -> bool: + lowered = str(message or "").lower() + return ( + "tool_call" in lowered + or "tool calls" in lowered + or "valid json" in lowered + or "unsupported tool" in lowered + or "invalid action" in lowered + ) + + # --------------------------------------------------------------------------- # Frozen-lake-specific tool-call parser (passed to generic client) # --------------------------------------------------------------------------- @@ -224,6 +770,8 @@ def __init__( allow_plaintext_action_fallback: bool = False, logprobs: bool = True, enable_thinking: Optional[bool] = False, + observation_mode: str = "text", + max_parse_retries: int = 0, ): self.model_id = model_id self.tokenizer_name_or_path = tokenizer_name_or_path @@ -237,6 +785,8 @@ def __init__( self.allow_plaintext_action_fallback = allow_plaintext_action_fallback self.logprobs = logprobs self.enable_thinking = enable_thinking + self.observation_mode = observation_mode + self.max_parse_retries = max_parse_retries def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> List[asyncio.Task[EvaluationRow]]: completion_params = dict(config.completion_params or {}) @@ -256,6 +806,8 @@ def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> or self.tokenizer_name_or_path or model_id ) + if tokenizer_name_or_path == model_id and _is_kimi_tokenizer_name(model_id): + tokenizer_name_or_path = "moonshotai/Kimi-K2.5" api_key = completion_params.get("api_key") or processor_kwargs.get("api_key") or self.api_key base_url = completion_params.get("base_url") or processor_kwargs.get("base_url") or self.base_url temperature = float(completion_params.get("temperature", processor_kwargs.get("temperature", self.temperature))) @@ -274,21 +826,35 @@ def __call__(self, rows: List[EvaluationRow], config: RolloutProcessorConfig) -> ), ) ) + observation_mode = str( + completion_params.get("observation_mode") + or processor_kwargs.get("observation_mode") + or self.observation_mode + ).strip().lower() + max_parse_retries = int( + completion_params.get("max_parse_retries") + or processor_kwargs.get("max_parse_retries") + or self.max_parse_retries + ) + if observation_mode == "image": + request_params.setdefault("thinking", {"type": "disabled"}) + fallback_system_prompt = DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS default_system_prompt = ( completion_params.get("system_prompt") or processor_kwargs.get("system_prompt") or self.system_prompt - or ( - "/no_think\n" - "You are an RL policy for FrozenLake.\n" - "Pick the action that moves toward G while avoiding H.\n" - "Always respond with exactly one tool call, no text." - ) + or fallback_system_prompt ) default_user_prompt_template = ( processor_kwargs.get("user_prompt_template") or self.user_prompt_template or completion_params.get("user_prompt_template") + or None + ) + default_visual_prompt = ( + processor_kwargs.get("visual_prompt_template") + or completion_params.get("visual_prompt_template") + or DEFAULT_VISUAL_PROMPT_TEMPLATE ) async def process_row(row: EvaluationRow) -> EvaluationRow: @@ -296,21 +862,39 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: tool_call_parser = build_frozen_lake_tool_call_parser( allow_plaintext_action_fallback=allow_plaintext_action_fallback, + tokenizer_getter=lambda: _get_hf_tokenizer(str(tokenizer_name_or_path)), model_id=model_id, ) - client = FireworksV1CompletionsClient( - model_id=model_id, - tokenizer_name_or_path=str(tokenizer_name_or_path), - api_key=str(api_key) if api_key is not None else None, - base_url=str(base_url) if base_url is not None else None, - temperature=temperature, - max_tokens=max_tokens, - request_params=request_params, - logprobs=self.logprobs, - enable_thinking=self.enable_thinking, - tool_call_parser=tool_call_parser, - default_tools=FROZEN_LAKE_TOOLS, - ) + text_client = None + image_client = None + if observation_mode == "image": + image_client = FireworksV1ImageCompletionsClient( + model_id=model_id, + tokenizer_name_or_path=str(tokenizer_name_or_path), + api_key=str(api_key) if api_key is not None else None, + base_url=str(base_url) if base_url is not None else None, + temperature=temperature, + max_tokens=max_tokens, + request_params=request_params, + logprobs=self.logprobs, + enable_thinking=self.enable_thinking, + tool_call_parser=tool_call_parser, + default_tools=FROZEN_LAKE_TOOLS, + ) + else: + text_client = FireworksV1CompletionsClient( + model_id=model_id, + tokenizer_name_or_path=str(tokenizer_name_or_path), + api_key=str(api_key) if api_key is not None else None, + base_url=str(base_url) if base_url is not None else None, + temperature=temperature, + max_tokens=max_tokens, + request_params=request_params, + logprobs=self.logprobs, + enable_thinking=self.enable_thinking, + tool_call_parser=tool_call_parser, + default_tools=FROZEN_LAKE_TOOLS, + ) usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} all_prompt_ids: List[int] = [] all_completion_ids: List[int] = [] @@ -321,6 +905,15 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: latest_finish_reason: Optional[str] = None latest_raw_output: Optional[Dict[str, Any]] = None rollout_error: Optional[str] = None + assistant_toolcall_prefill_ids: List[int] = [] + assistant_toolcall_prefill_text = "" + if _is_kimi_tokenizer_name(str(tokenizer_name_or_path)): + assistant_toolcall_prefill_ids = _build_kimi_toolcall_generation_prefill_token_ids( + str(tokenizer_name_or_path) + ) + assistant_toolcall_prefill_text = _build_kimi_toolcall_generation_prefill_text( + str(tokenizer_name_or_path) + ) try: dataset_info = row.input_metadata.dataset_info or {} @@ -332,6 +925,9 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: row_user_prompt_template = dataset_info.get("user_prompt_template") if not isinstance(row_user_prompt_template, str) or not row_user_prompt_template: row_user_prompt_template = default_user_prompt_template + row_visual_prompt_template = dataset_info.get("visual_prompt_template") + if not isinstance(row_visual_prompt_template, str) or not row_visual_prompt_template: + row_visual_prompt_template = default_visual_prompt env = build_frozen_lake_tool_env(environment_context=env_context, max_steps=max_steps) state = env.reset() @@ -342,15 +938,45 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: if default_system_prompt and not any(m.role == "system" for m in messages): messages.insert(0, Message(role="system", content=default_system_prompt)) row.tools = row.tools or list(FROZEN_LAKE_TOOLS) - initial_user_prompt = build_frozen_lake_user_prompt( - user_prompt_template=str(row_user_prompt_template) if row_user_prompt_template else None, - observation=observation, - ) - messages.append(Message(role="user", content=initial_user_prompt)) - current_prompt_ids = client.build_prompt_token_ids( - messages=_to_message_payload(messages=messages), - tools=row.tools, - ) + current_prompt_ids: List[int] = [] + current_images: List[str] = [] + current_prompt_text = "" + if observation_mode != "image": + initial_user_prompt = build_frozen_lake_user_prompt( + user_prompt_template=str(row_user_prompt_template) if row_user_prompt_template else None, + observation=observation, + ) + messages.append(Message(role="user", content=initial_user_prompt)) + current_prompt_ids = text_client.build_prompt_token_ids( + messages=_to_message_payload(messages=messages), + tools=row.tools, + ) + else: + initial_visual_prompt = _build_visual_user_prompt( + prompt_template=str(row_visual_prompt_template) if row_visual_prompt_template else None, + observation=observation, + default_prompt=str(default_visual_prompt), + ) + initial_image_url = env.render_image_data_url() + current_images = [initial_image_url] + messages.append( + Message( + role="user", + content=_build_multimodal_image_content( + image_url=initial_image_url, + text=initial_visual_prompt, + ), + ) + ) + message_payloads = _to_message_payload(messages=messages) + current_prompt_ids = image_client.build_prompt_token_ids( + messages=message_payloads, + tools=row.tools, + ) + current_prompt_text = image_client.build_prompt_text( + messages=message_payloads, + tools=row.tools, + ) for step_index in range(max_steps): if done: @@ -363,22 +989,59 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: except Exception: pass + base_prompt_ids = list(current_prompt_ids) + base_prompt_text = str(current_prompt_text) + model_request_traces.append( { "step_index": step_index + 1, - "prompt_ids": list(current_prompt_ids), - "prompt_token_count": len(current_prompt_ids), + "prompt_ids": list(base_prompt_ids), + "prompt_token_count": len(base_prompt_ids), "tools": list(row.tools or []), + "observation_mode": observation_mode, } ) - try: - completion = await client.create_completion_from_prompt_ids( - prompt_token_ids=current_prompt_ids, - tools=row.tools, - ) - except Exception as exc: - rollout_error = str(exc) - tool_call_traces.append({"step_index": step_index + 1, "error": rollout_error}) + request_prompt_ids = list(base_prompt_ids) + request_prompt_text = base_prompt_text + if assistant_toolcall_prefill_ids: + request_prompt_ids.extend(assistant_toolcall_prefill_ids) + if observation_mode == "image": + request_prompt_text = f"{request_prompt_text}{assistant_toolcall_prefill_text}" + model_request_traces[-1]["assistant_prefill_len"] = len(assistant_toolcall_prefill_ids) + model_request_traces[-1]["assistant_prefill_text"] = assistant_toolcall_prefill_text + model_request_traces[-1]["logical_prompt_ids"] = list(base_prompt_ids) + model_request_traces[-1]["logical_prompt_token_count"] = len(base_prompt_ids) + model_request_traces[-1]["prompt_ids"] = list(request_prompt_ids) + model_request_traces[-1]["prompt_token_count"] = len(request_prompt_ids) + parse_retry_count = 0 + while True: + try: + if observation_mode == "image": + completion = await image_client.create_completion_from_prompt_ids( + prompt_token_ids=request_prompt_ids, + prompt_text=request_prompt_text, + images=current_images, + tools=row.tools, + ) + model_request_traces[-1]["image_count"] = len(current_images) + model_request_traces[-1]["prompt_text"] = request_prompt_text + else: + completion = await text_client.create_completion_from_prompt_ids( + prompt_token_ids=request_prompt_ids, + tools=row.tools, + ) + break + except Exception as exc: + if parse_retry_count >= max_parse_retries or not _looks_like_tool_parse_error_message(str(exc)): + rollout_error = str(exc) + tool_call_traces.append({"step_index": step_index + 1, "error": rollout_error}) + completion = None + break + parse_retry_count += 1 + model_request_traces[-1]["parse_retry_count"] = parse_retry_count + model_request_traces[-1].setdefault("parse_retry_errors", []).append(str(exc)) + await asyncio.sleep(0.25 * parse_retry_count) + if completion is None: break usage_payload = completion.get("usage") or {} @@ -386,10 +1049,23 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: usage["completion_tokens"] += int(usage_payload.get("completion_tokens", 0)) usage["total_tokens"] += int(usage_payload.get("total_tokens", 0)) - prompt_ids = [int(x) for x in list(completion.get("prompt_ids") or current_prompt_ids)] - completion_ids = [int(x) for x in list(completion.get("completion_ids") or [])] + provider_prompt_ids = [int(x) for x in list(completion.get("prompt_ids") or [])] + prompt_ids = [int(x) for x in list(base_prompt_ids)] + raw_completion_ids = [int(x) for x in list(completion.get("completion_ids") or [])] + if observation_mode == "image": + assistant_suffix_ids = image_client.encode_assistant_turn_suffix() + else: + assistant_suffix_ids = text_client.encode_special_suffix() + completion_ids = ( + [int(x) for x in list(assistant_toolcall_prefill_ids)] + + list(raw_completion_ids) + + [int(x) for x in list(assistant_suffix_ids)] + ) + completion_text = str(completion.get("completion_text") or "") all_prompt_ids.extend(prompt_ids) all_completion_ids.extend(completion_ids) + if provider_prompt_ids and provider_prompt_ids != request_prompt_ids: + model_request_traces[-1]["provider_prompt_ids"] = provider_prompt_ids choice = (completion.get("choices") or [{}])[0] message_payload = choice.get("message") or {} @@ -400,7 +1076,7 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: assistant_message_payload = { "role": "assistant", - "content": message_payload.get("content", ""), + "content": completion_text or str(message_payload.get("content", "") or ""), "tool_calls": message_payload.get("tool_calls"), } assistant_message = Message.model_validate(assistant_message_payload) @@ -426,13 +1102,15 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: raw_output.get("tool_call_parser") if isinstance(raw_output, dict) else None ), } + if observation_mode == "image": + turn_trace["prompt_text"] = base_prompt_text + turn_trace["observation_mode"] = "image" completion_logprobs = completion.get("completion_logprobs") if completion_logprobs: turn_trace["completion_logprobs"] = [float(lp) for lp in completion_logprobs] token_turn_traces.append(turn_trace) tool_result = { - "observation": observation, "action": action, "reward": reward, "terminated": bool(state.get("terminated", False)), @@ -443,22 +1121,59 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: "tile": state.get("tile"), "step_index": state.get("step_index", step_index + 1), } + tool_image_url = env.render_image_data_url() if observation_mode == "image" else None tool_message_payload = { "role": "tool", "name": TOOL_NAME_LAKE_MOVE, "tool_call_id": tool_call_id or None, - "content": json.dumps(tool_result, separators=(",", ":")), + "content": ( + _build_multimodal_image_content(image_url=tool_image_url) + if observation_mode == "image" and tool_image_url + else json.dumps(tool_result, separators=(",", ":")) + ), } messages.append(Message.model_validate(tool_message_payload)) - im_end_ids = client.encode_special_suffix() - assistant_turn_ids = list(completion_ids) + im_end_ids - tool_suffix_ids = client.build_tool_response_suffix_token_ids( - tool_message=tool_message_payload - ) - current_prompt_ids = list(prompt_ids) + list(assistant_turn_ids) + list(tool_suffix_ids) - model_request_traces[-1]["assistant_turn_len"] = len(completion_ids) - model_request_traces[-1]["tool_suffix_len"] = len(tool_suffix_ids) + if observation_mode == "image": + model_request_traces[-1]["assistant_turn_len"] = len(completion_ids) + if not done: + assistant_turn_ids = list(completion_ids) + assistant_turn_text = ( + str(assistant_toolcall_prefill_text or "") + + str(completion_text or "") + + image_client.assistant_turn_suffix_text() + ) + tool_suffix_ids = image_client.build_tool_response_suffix_token_ids( + tool_message=tool_message_payload, + ) + tool_suffix_text = image_client.build_tool_response_suffix_text( + tool_message=tool_message_payload, + ) + tool_suffix_text = image_client.decode_token_ids( + token_ids=tool_suffix_ids + ) + current_prompt_ids = ( + list(prompt_ids) + list(assistant_turn_ids) + list(tool_suffix_ids) + ) + current_prompt_text = ( + str(base_prompt_text) + + assistant_turn_text + + tool_suffix_text + ) + if tool_image_url is not None: + current_images = list(current_images) + [tool_image_url] + model_request_traces[-1]["tool_suffix_len"] = len(tool_suffix_ids) + else: + current_prompt_ids = [] + current_prompt_text = "" + else: + assistant_turn_ids = list(completion_ids) + tool_suffix_ids = text_client.build_tool_response_suffix_token_ids( + tool_message=tool_message_payload + ) + current_prompt_ids = list(prompt_ids) + list(assistant_turn_ids) + list(tool_suffix_ids) + model_request_traces[-1]["assistant_turn_len"] = len(completion_ids) + model_request_traces[-1]["tool_suffix_len"] = len(tool_suffix_ids) tool_call_traces.append( { "step_index": step_index + 1, @@ -493,6 +1208,10 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: extra["tools_input"] = list(row.tools or []) extra["model_request_traces"] = list(model_request_traces) extra["token_turn_traces"] = list(token_turn_traces) + extra["observation_mode"] = observation_mode + extra["tool_call_generation_mode"] = ( + "plaintext_fallback" if allow_plaintext_action_fallback else "prompt_only" + ) if rollout_error: extra["rollout_error"] = rollout_error row.execution_metadata.extra = extra @@ -504,7 +1223,10 @@ async def process_row(row: EvaluationRow) -> EvaluationRow: pass return row finally: - await client.close() + if text_client is not None: + await text_client.close() + if image_client is not None: + await image_client.close() async def _sem_wrapper(target_row: EvaluationRow) -> EvaluationRow: async with semaphore: diff --git a/training/examples/frozen_lake/frozen_lake_schema.py b/training/examples/frozen_lake/frozen_lake_schema.py index 60c28645..15d49491 100644 --- a/training/examples/frozen_lake/frozen_lake_schema.py +++ b/training/examples/frozen_lake/frozen_lake_schema.py @@ -111,9 +111,53 @@ def _load_xml_tool_call_payload_with_text_span(text: str) -> Tuple[Dict[str, Any return payload, start, end +def _load_kimi_native_tool_call_payload_with_text_span(text: str) -> Tuple[Dict[str, Any], int, int]: + match = re.search( + r"(?:<\|tool_calls_section_begin\|>)?\s*" + r"<\|tool_call_begin\|>\s*(.*?)\s*" + r"<\|tool_call_argument_begin\|>\s*(\{.*?\})\s*" + r"<\|tool_call_end\|>\s*" + r"(?:<\|tool_calls_section_end\|>)?", + text, + flags=re.DOTALL, + ) + if not match: + raise ValueError("No Kimi native tool_call block found") + + raw_identifier = str(match.group(1) or "").strip() + arguments_text = match.group(2) + try: + arguments = json.loads(arguments_text) + except json.JSONDecodeError as exc: + raise ValueError(f"Invalid JSON inside Kimi tool call: {arguments_text!r}") from exc + if not isinstance(arguments, dict): + raise ValueError("Kimi tool call arguments must be a JSON object") + + tool_name = TOOL_NAME_LAKE_MOVE + identifier = raw_identifier + if raw_identifier.startswith("functions."): + qualified_name = raw_identifier.split(":", 1)[0] + tool_name = qualified_name.split(".", 1)[1] or TOOL_NAME_LAKE_MOVE + payload = { + "tool_calls": [ + { + "id": identifier or f"toolcall_{uuid.uuid4().hex[:12]}", + "name": tool_name, + "arguments": arguments, + } + ] + } + start, end = match.span() + return payload, start, end + + def parse_first_frozen_lake_tool_call(output_text: str) -> ParsedToolCall: """Parse the first tool call from strict tool-call JSON output.""" - loaders = (_load_json_object_with_text_span, _load_xml_tool_call_payload_with_text_span) + loaders = ( + _load_json_object_with_text_span, + _load_xml_tool_call_payload_with_text_span, + _load_kimi_native_tool_call_payload_with_text_span, + ) last_error: Exception | None = None for loader in loaders: try: @@ -129,7 +173,11 @@ def parse_first_frozen_lake_tool_call(output_text: str) -> ParsedToolCall: def parse_first_frozen_lake_tool_call_with_content(output_text: str) -> Tuple[ParsedToolCall, str]: """Parse the first tool call and return residual assistant text content.""" - loaders = (_load_json_object_with_text_span, _load_xml_tool_call_payload_with_text_span) + loaders = ( + _load_json_object_with_text_span, + _load_xml_tool_call_payload_with_text_span, + _load_kimi_native_tool_call_payload_with_text_span, + ) last_error: Exception | None = None for loader in loaders: try: diff --git a/training/examples/frozen_lake/rendering.py b/training/examples/frozen_lake/rendering.py new file mode 100644 index 00000000..1b81522e --- /dev/null +++ b/training/examples/frozen_lake/rendering.py @@ -0,0 +1,119 @@ +"""Rendering helpers for FrozenLake visual rollouts. + +This mirrors the upstream Gymnasium FrozenLake renderer, but uses PIL so we can +produce deterministic PNG data URLs without depending on pygame at rollout time. +""" + +from __future__ import annotations + +import base64 +import io +from functools import lru_cache +from pathlib import Path +from typing import Sequence + +from PIL import Image, ImageDraw + + +ASSET_DIR = Path(__file__).resolve().parent / "assets" / "img" +DEFAULT_ACTION = "DOWN" +ASSET_BY_TILE = { + "H": "hole.png", + "G": "goal.png", + "S": "stool.png", +} +ELF_ASSET_BY_ACTION = { + "LEFT": "elf_left.png", + "DOWN": "elf_down.png", + "RIGHT": "elf_right.png", + "UP": "elf_up.png", +} +try: + RESAMPLE_LANCZOS = Image.Resampling.LANCZOS +except AttributeError: + RESAMPLE_LANCZOS = Image.LANCZOS + + +@lru_cache(maxsize=64) +def _load_scaled_asset(asset_name: str, cell_size: int) -> Image.Image: + asset_path = ASSET_DIR / asset_name + image = Image.open(asset_path).convert("RGBA") + return image.resize((cell_size, cell_size), RESAMPLE_LANCZOS) + + +def render_frozen_lake_png_bytes( + map_rows: Sequence[str], + *, + agent_row: int, + agent_col: int, + last_action: str | None = None, + cell_size: int = 96, +) -> bytes: + rows = len(map_rows) + cols = len(map_rows[0]) if map_rows else 0 + canvas = Image.new("RGBA", (cols * cell_size, rows * cell_size), (255, 255, 255, 255)) + draw = ImageDraw.Draw(canvas) + + ice_tile = _load_scaled_asset("ice.png", cell_size) + hole_tile = _load_scaled_asset("hole.png", cell_size) + cracked_hole_tile = _load_scaled_asset("cracked_hole.png", cell_size) + goal_tile = _load_scaled_asset("goal.png", cell_size) + start_tile = _load_scaled_asset("stool.png", cell_size) + elf_tile = _load_scaled_asset( + ELF_ASSET_BY_ACTION.get(str(last_action or DEFAULT_ACTION).upper(), "elf_down.png"), + cell_size, + ) + + overlay_by_tile = { + "H": hole_tile, + "G": goal_tile, + "S": start_tile, + } + + for row_index, row in enumerate(map_rows): + for col_index, cell in enumerate(row): + position = (col_index * cell_size, row_index * cell_size) + canvas.alpha_composite(ice_tile, position) + overlay = overlay_by_tile.get(cell) + if overlay is not None: + canvas.alpha_composite(overlay, position) + draw.rectangle( + ( + position[0], + position[1], + position[0] + cell_size, + position[1] + cell_size, + ), + outline=(180, 200, 230, 255), + width=1, + ) + + agent_position = (agent_col * cell_size, agent_row * cell_size) + agent_tile = map_rows[agent_row][agent_col] + if agent_tile == "H": + canvas.alpha_composite(cracked_hole_tile, agent_position) + else: + canvas.alpha_composite(elf_tile, agent_position) + + out = io.BytesIO() + canvas.convert("RGB").save(out, format="PNG") + return out.getvalue() + + +def render_frozen_lake_png_data_url( + map_rows: Sequence[str], + *, + agent_row: int, + agent_col: int, + last_action: str | None = None, + cell_size: int = 96, +) -> str: + png_bytes = render_frozen_lake_png_bytes( + map_rows, + agent_row=agent_row, + agent_col=agent_col, + last_action=last_action, + cell_size=cell_size, + ) + encoded = base64.b64encode(png_bytes).decode("ascii") + return f"data:image/png;base64,{encoded}" diff --git a/training/examples/frozen_lake/test_masking.py b/training/examples/frozen_lake/test_masking.py index ca06f393..47185446 100644 --- a/training/examples/frozen_lake/test_masking.py +++ b/training/examples/frozen_lake/test_masking.py @@ -18,6 +18,7 @@ build_ui_token_mask, compute_model_output_spans, ) +from training.utils.supervised import build_datum_from_token_mask def _make_traces(turns: list[dict]) -> tuple[list[dict], list[dict]]: @@ -78,6 +79,18 @@ def _make_traces(turns: list[dict]) -> tuple[list[dict], list[dict]]: }, ] +KIMI_STYLE_MULTI_TURN = [ + { + "prompt_ids": [1, 2, 3], + "completion_ids": [99, 10, 90], # prefill, raw tool call token(s), im_end + "assistant_turn_len": 3, + }, + { + "prompt_ids": [1, 2, 3, 99, 10, 90, 20, 21], + "completion_ids": [99, 30, 90], + }, +] + @pytest.mark.parametrize("scenario", [SINGLE_TURN, MULTI_TURN_2, MULTI_TURN_3], ids=["single", "multi_2", "multi_3"]) @@ -183,3 +196,32 @@ def test_fallback_when_assistant_turn_len_missing(): assert spans_no_mrt[0] == (3, 2, 1), "should fall back to completion_ids len" assert spans_with_mrt[0] == spans_no_mrt[0] + + +def test_kimi_style_mask_includes_prefill_and_im_end_tokens(): + """Assistant framing tokens should be masked as model output too.""" + ttt, mrt = _make_traces(KIMI_STYLE_MULTI_TURN) + spans = compute_model_output_spans(ttt, mrt) + + assert spans == [(3, 3, 1), (8, 3, 2)] + + ui_mask = build_ui_token_mask(spans, 11) + assert ui_mask == [0, 0, 0, 1, 1, 1, 0, 0, 2, 2, 2] + + train_mask = build_training_loss_mask(spans, 10) + assert train_mask == [0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0] + + +def test_shared_weighted_datum_matches_training_loss_mask(): + """The shared datum helper should use the exact same shifted mask as training.""" + ttt, mrt = _make_traces(KIMI_STYLE_MULTI_TURN) + spans = compute_model_output_spans(ttt, mrt) + + full_ids = list(ttt[-1]["prompt_ids"]) + list(ttt[-1]["completion_ids"]) + ui_mask = build_ui_token_mask(spans, len(full_ids)) + train_mask = build_training_loss_mask(spans, len(full_ids) - 1) + + rendered = build_datum_from_token_mask(full_ids, ui_mask, include_loss_mask=True) + + assert rendered.datum.loss_fn_inputs["weights"].data == train_mask + assert rendered.datum.loss_fn_inputs["loss_mask"].data == train_mask diff --git a/training/examples/frozen_lake/train_frozen_lake.py b/training/examples/frozen_lake/train_frozen_lake.py index 4b793530..9dcd023c 100644 --- a/training/examples/frozen_lake/train_frozen_lake.py +++ b/training/examples/frozen_lake/train_frozen_lake.py @@ -36,10 +36,13 @@ from eval_protocol.pytest.types import RolloutProcessorConfig from training.examples.frozen_lake.frozen_lake_env import build_frozen_lake_tool_env -from training.examples.frozen_lake.frozen_lake_rollout import FrozenLakeToolRolloutProcessor +from training.examples.frozen_lake.frozen_lake_rollout import ( + DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS, + FrozenLakeToolRolloutProcessor, +) from training.examples.frozen_lake.masking import ( compute_model_output_spans, - build_training_loss_mask, + build_ui_token_mask, ) from fireworks.training.sdk import DeploymentManager, TrainerJobManager @@ -61,6 +64,7 @@ setup_deployment, create_trainer_job, compute_advantages, + build_datum_from_token_mask, ) from training.utils.rl import PromptGroup from training.utils.rl.train import MinibatchTrainFns, run_rl_loop @@ -82,14 +86,16 @@ # --------------------------------------------------------------------------- DEFAULT_USER_PROMPT_TEMPLATE = "FrozenLake grid observation:\n{observation}" - -DEFAULT_SYSTEM_PROMPT = ( - "/no_think\n" - "You are an RL policy for FrozenLake.\n" - "Pick the action that moves toward G while avoiding H.\n" - "Always respond with exactly one tool call, no text." +DEFAULT_VISUAL_PROMPT_TEMPLATE = ( + "You are playing FrozenLake. The image shows the current grid. " + "The current textual observation is below.\n" + "{observation}\n\n" + "Tiles are labeled S, F, H, and G. The agent is marked with a red dot. " + "Use exactly one lake_move tool call with action LEFT, DOWN, RIGHT, or UP." ) +DEFAULT_SYSTEM_PROMPT = DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS + @dataclass class FrozenLakeConfig: @@ -121,6 +127,9 @@ class FrozenLakeConfig: use_random_map: bool = True system_prompt: str = DEFAULT_SYSTEM_PROMPT user_prompt_template: str = DEFAULT_USER_PROMPT_TEMPLATE + visual_prompt_template: str = DEFAULT_VISUAL_PROMPT_TEMPLATE + observation_mode: str = "text" + allow_plaintext_action_fallback: bool = False training_shape: str = "" deployment_shape: str = "" @@ -166,6 +175,9 @@ def parse_args() -> FrozenLakeConfig: parser.add_argument("--wandb-entity", default=os.environ.get("WANDB_ENTITY", "")) parser.add_argument("--wandb-project", default=os.environ.get("WANDB_PROJECT", "frozen-lake-grpo")) + parser.add_argument("--visual-prompt-template", default=DEFAULT_VISUAL_PROMPT_TEMPLATE) + parser.add_argument("--observation-mode", choices=("text", "image"), default="text") + parser.add_argument("--allow-plaintext-action-fallback", action="store_true") return cast(FrozenLakeConfig, parser.parse_args(namespace=FrozenLakeConfig())) @@ -234,23 +246,16 @@ def evaluation_row_to_training_data( # prompt_len = first turn's prompt length (system + user message, before any model output) first_prompt_len = len([int(x) for x in (token_turn_traces[0].get("prompt_ids") or [])]) - model_input_len = len(full_tokens) - 1 - model_request_traces = extra.get("model_request_traces") or [] spans = compute_model_output_spans(token_turn_traces, model_request_traces) - loss_mask = build_training_loss_mask(spans, model_input_len) - - datum = tinker.Datum( - model_input=tinker.ModelInput.from_ints(full_tokens[:-1]), - loss_fn_inputs={ - "target_tokens": tinker.TensorData( - data=full_tokens[1:], dtype="int64", shape=[model_input_len] - ), - "loss_mask": tinker.TensorData( - data=loss_mask, dtype="float32", shape=[model_input_len] - ), - }, + token_mask = build_ui_token_mask(spans, len(full_tokens)) + rendered = build_datum_from_token_mask( + full_tokens, + token_mask, + include_loss_mask=True, ) + datum = rendered.datum + model_input_len = len(rendered.token_ids) - 1 # Reconstruct inference logprobs aligned to the full sequence. # Pad prompt positions with 0.0, then fill in per-turn completion logprobs. @@ -559,6 +564,8 @@ def _cleanup_deployment(deployment_id: str | None) -> None: system_prompt=cfg.system_prompt, user_prompt_template=cfg.user_prompt_template, logprobs=True, + observation_mode=cfg.observation_mode, + allow_plaintext_action_fallback=cfg.allow_plaintext_action_fallback, ) rollout_config = RolloutProcessorConfig( completion_params={"model": inference_model}, @@ -590,6 +597,7 @@ async def sample_one_prompt(env_context: Dict[str, Any]) -> PromptGroup | None: dataset_info={ "environment_context": dict(env_context), "user_prompt_template": cfg.user_prompt_template, + "visual_prompt_template": cfg.visual_prompt_template, }, ), )) diff --git a/training/examples/frozen_lake/validation_screenshots/kimi-ep-logs.png b/training/examples/frozen_lake/validation_screenshots/kimi-ep-logs.png new file mode 100644 index 00000000..6bd6987f Binary files /dev/null and b/training/examples/frozen_lake/validation_screenshots/kimi-ep-logs.png differ diff --git a/training/examples/frozen_lake/validation_screenshots/qwen3-vl-ep-logs.png b/training/examples/frozen_lake/validation_screenshots/qwen3-vl-ep-logs.png new file mode 100644 index 00000000..6a188436 Binary files /dev/null and b/training/examples/frozen_lake/validation_screenshots/qwen3-vl-ep-logs.png differ diff --git a/training/examples/frozen_lake/verify_rollout.py b/training/examples/frozen_lake/verify_rollout.py index 40b9d4ba..a88472e0 100644 --- a/training/examples/frozen_lake/verify_rollout.py +++ b/training/examples/frozen_lake/verify_rollout.py @@ -1,8 +1,8 @@ #!/usr/bin/env python3 -"""Run a few rollouts on qwen3-8b serverless and write eval-protocol JSONL for visualization. +"""Run a few FrozenLake rollouts and write eval-protocol rows for inspection. Usage: - python verify_rollout.py [--num-rollouts 4] [--port 8765] + python verify_rollout.py [--num-rollouts 4] [--port 8765] [--model-id ...] [--visual] After it finishes, the vite dashboard opens automatically at http://localhost:/table with token-level debug data in execution_metadata.extra. @@ -12,11 +12,14 @@ import argparse import asyncio +import html import json import os import sys import time import logging +import threading +import webbrowser from pathlib import Path logging.basicConfig(level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S") @@ -30,53 +33,97 @@ from eval_protocol.pytest.types import RolloutProcessorConfig from eval_protocol.dataset_logger import default_logger -from training.examples.frozen_lake.frozen_lake_rollout import FrozenLakeToolRolloutProcessor +from training.examples.frozen_lake.frozen_lake_rollout import ( + DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS, + FrozenLakeToolRolloutProcessor, +) +from training.examples.frozen_lake.frozen_lake_env import build_frozen_lake_tool_env from training.examples.frozen_lake.masking import compute_model_output_spans, build_ui_token_mask -DEFAULT_SYSTEM_PROMPT = ( - "/no_think\n" - "You are an RL policy for FrozenLake.\n" - "Pick the action that moves toward G while avoiding H.\n" - "Always respond with exactly one tool call, no text." -) +DEFAULT_LOGS_DB_PATH = str(COOKBOOK_DIR / ".tmp" / "frozen_lake_verify_logs.db") +DEFAULT_REPORT_PATH = str(COOKBOOK_DIR / ".tmp" / "frozen_lake_verify_report.html") +DEFAULT_KIMI_TOKENIZER = "moonshotai/Kimi-K2.5" +DEFAULT_SYSTEM_PROMPT = DEFAULT_SYSTEM_PROMPT_INSTRUCTIONS -def load_seed_contexts(path: str, max_seeds: int = 4): + +def load_seed_contexts(path: str, max_seeds: int = 4, seed_values: list[int] | None = None): + wanted_seeds = {int(seed) for seed in (seed_values or [])} contexts = [] with open(path) as f: for line in f: ctx = json.loads(line) + if wanted_seeds and int(ctx.get("seed", -1)) not in wanted_seeds: + continue contexts.append(ctx) if len(contexts) >= max_seeds: break return contexts -async def run_rollouts(num_rollouts: int = 4): +def _looks_like_kimi_target(value: object) -> bool: + lowered = str(value or "").strip().lower() + return "kimi-k2.5" in lowered or "kimi-k2p5" in lowered + + +def _resolve_tokenizer_model(model_id: str, tokenizer_model: str | None) -> str: + if tokenizer_model: + return str(tokenizer_model) + env_tokenizer = os.environ.get("FROZEN_LAKE_VERIFY_TOKENIZER_MODEL", "").strip() + if env_tokenizer: + return env_tokenizer + if _looks_like_kimi_target(model_id): + return DEFAULT_KIMI_TOKENIZER + return "Qwen/Qwen3-8B" + + +def _validation_target(*, model_id: str, tokenizer_name: str, visual: bool) -> str: + if visual and (_looks_like_kimi_target(model_id) or _looks_like_kimi_target(tokenizer_name)): + return "kimi-k2.5-vl" + if visual: + return "generic-vlm" + if _looks_like_kimi_target(model_id) or _looks_like_kimi_target(tokenizer_name): + return "kimi-k2.5-text" + return "text" + + +async def run_rollouts( + *, + num_rollouts: int = 4, + rows_per_seed: int = 1, + model_id: str = "accounts/fireworks/models/qwen3-8b", + tokenizer: str | None = None, + visual: bool = False, + allow_plaintext_action_fallback: bool = False, + temperature: float = 0.0, + seed_values: list[int] | None = None, +): api_key = os.environ["FIREWORKS_API_KEY"] - model_id = "accounts/fireworks/models/qwen3-8b" - tokenizer = "Qwen/Qwen3-8B" + resolved_tokenizer = _resolve_tokenizer_model(model_id=model_id, tokenizer_model=tokenizer) seed_path = str(FROZEN_LAKE_DIR / "seeds.jsonl") - seed_contexts = load_seed_contexts(seed_path, max_seeds=num_rollouts) + seed_contexts = load_seed_contexts(seed_path, max_seeds=num_rollouts, seed_values=seed_values) logger.info("Loaded %d seeds", len(seed_contexts)) processor = FrozenLakeToolRolloutProcessor( model_id=model_id, - tokenizer_name_or_path=tokenizer, + tokenizer_name_or_path=resolved_tokenizer, api_key=api_key, base_url="https://api.fireworks.ai/inference", - temperature=1.0, + temperature=float(temperature), max_tokens=128, system_prompt=DEFAULT_SYSTEM_PROMPT, logprobs=True, + observation_mode="image" if visual else "text", + allow_plaintext_action_fallback=allow_plaintext_action_fallback, + max_parse_retries=3 if visual and _looks_like_kimi_target(model_id) else 0, ) config = RolloutProcessorConfig( - completion_params={"model": model_id}, + completion_params={"model": model_id, "temperature": float(temperature)}, mcp_config_path="", steps=30, - semaphore=asyncio.Semaphore(8), + semaphore=asyncio.Semaphore(1), ) all_results = [] @@ -88,7 +135,7 @@ async def run_rollouts(num_rollouts: int = 4): dataset_info={"environment_context": dict(ctx)}, ), ) - for i in range(2) + for i in range(rows_per_seed) ] tasks = processor(rows, config) @@ -106,10 +153,230 @@ async def run_rollouts(num_rollouts: int = 4): return all_results -def enrich_rows(results: list[EvaluationRow], tokenizer_name: str): +def _build_row_validation_details( + row: EvaluationRow, + *, + model_id: str, + tokenizer_name: str, + visual: bool, +) -> dict[str, object]: + extra = row.execution_metadata.extra if isinstance(row.execution_metadata.extra, dict) else {} + token_turn_traces = extra.get("token_turn_traces") or [] + model_request_traces = extra.get("model_request_traces") or [] + tool_call_traces = extra.get("tool_call_traces") or [] + full_episode = extra.get("full_episode") or {} + prompt_lengths = [len(trace.get("prompt_ids") or []) for trace in token_turn_traces] + completion_lengths = [len(trace.get("completion_ids") or []) for trace in token_turn_traces] + image_counts = [trace.get("image_count") for trace in model_request_traces] + token_ids = list(full_episode.get("token_ids") or []) + mask = list(full_episode.get("mask") or []) + logprobs = list(full_episode.get("logprobs") or []) + + validation_checks: list[dict[str, str]] = [] + failing_checks: list[str] = [] + skipped_checks: list[str] = [] + + def add_check(name: str, status: str, detail: str) -> None: + validation_checks.append({"name": name, "status": status, "detail": detail}) + if status == "fail": + failing_checks.append(name) + elif status == "skip": + skipped_checks.append(name) + + rollout_error = str(extra.get("rollout_error") or "").strip() + add_check( + "no_rollout_error", + "pass" if not rollout_error else "fail", + "row completed without rollout_error" if not rollout_error else rollout_error, + ) + + if visual: + observation_mode = str(extra.get("observation_mode") or "") + add_check( + "observation_mode_image", + "pass" if observation_mode == "image" else "fail", + f"observation_mode={observation_mode!r}", + ) + + add_check( + "token_turn_traces_present", + "pass" if token_turn_traces else "fail", + f"turn_count={len(token_turn_traces)}", + ) + add_check( + "prompt_ids_present_each_turn", + "pass" if token_turn_traces and all(trace.get("prompt_ids") for trace in token_turn_traces) else "fail", + f"prompt_lengths={prompt_lengths}", + ) + add_check( + "completion_ids_present_each_turn", + "pass" if token_turn_traces and all(trace.get("completion_ids") for trace in token_turn_traces) else "fail", + f"completion_lengths={completion_lengths}", + ) + + if len(token_turn_traces) < 2: + add_check("multi_turn_continuity", "skip", "requires at least two turns") + add_check("prompt_prefix_continuity", "skip", "requires at least two turns") + add_check("completion_replayed_in_next_prompt", "skip", "requires at least two turns") + else: + prompt_increase_ok = all(prompt_lengths[idx] > prompt_lengths[idx - 1] for idx in range(1, len(prompt_lengths))) + add_check( + "multi_turn_continuity", + "pass" if prompt_increase_ok else "fail", + f"prompt_lengths={prompt_lengths}", + ) + + prefix_ok = True + replay_ok = True + for idx in range(1, len(token_turn_traces)): + prev_prompt_ids = [int(x) for x in (token_turn_traces[idx - 1].get("prompt_ids") or [])] + prev_completion_ids = [int(x) for x in (token_turn_traces[idx - 1].get("completion_ids") or [])] + current_prompt_ids = [int(x) for x in (token_turn_traces[idx].get("prompt_ids") or [])] + if current_prompt_ids[: len(prev_prompt_ids)] != prev_prompt_ids: + prefix_ok = False + replay_start = len(prev_prompt_ids) + replay_end = replay_start + len(prev_completion_ids) + if current_prompt_ids[replay_start:replay_end] != prev_completion_ids: + replay_ok = False + + add_check( + "prompt_prefix_continuity", + "pass" if prefix_ok else "fail", + "next turn prompt_ids preserve the prior prompt_ids prefix", + ) + add_check( + "completion_replayed_in_next_prompt", + "pass" if replay_ok else "fail", + "next turn prompt_ids embed the prior completion_ids immediately after the prior prompt_ids", + ) + + request_count_matches = len(model_request_traces) == len(token_turn_traces) + add_check( + "model_request_trace_count_matches_turns", + "pass" if request_count_matches else "fail", + f"request_traces={len(model_request_traces)} turn_traces={len(token_turn_traces)}", + ) + add_check( + "tool_call_trace_count_matches_turns", + "pass" if len(tool_call_traces) == len(token_turn_traces) else "fail", + f"tool_call_traces={len(tool_call_traces)} turn_traces={len(token_turn_traces)}", + ) + + if visual: + image_count_ok = ( + bool(token_turn_traces) + and len(image_counts) == len(token_turn_traces) + and image_counts[0] == 1 + and all( + isinstance(image_counts[idx], int) + and image_counts[idx] == image_counts[idx - 1] + 1 + for idx in range(1, len(image_counts)) + ) + ) + add_check( + "image_counts_increment_per_turn", + "pass" if image_count_ok else "fail", + f"image_counts={image_counts}", + ) + prompt_text_present = bool(model_request_traces) and all( + str(trace.get("prompt_text") or "").strip() for trace in model_request_traces + ) + add_check( + "prompt_text_present_for_image_turns", + "pass" if prompt_text_present else "fail", + f"prompt_text_turns={sum(1 for trace in model_request_traces if str(trace.get('prompt_text') or '').strip())}", + ) + assistant_turn_len_ok = bool(model_request_traces) and all( + int(model_request_traces[idx].get("assistant_turn_len") or 0) == completion_lengths[idx] + for idx in range(min(len(model_request_traces), len(completion_lengths))) + ) + add_check( + "assistant_turn_len_matches_completion_ids", + "pass" if assistant_turn_len_ok else "fail", + "assistant_turn_len matches completion_ids length on every turn", + ) + if _looks_like_kimi_target(model_id) or _looks_like_kimi_target(tokenizer_name): + kimi_instant_prompt_ok = bool(token_turn_traces) and all( + "" in str(trace.get("prompt_text") or "") + for trace in token_turn_traces + ) + add_check( + "kimi_instant_prompt_contains_empty_think_block", + "pass" if kimi_instant_prompt_ok else "fail", + "each Kimi prompt_text should include when thinking is disabled", + ) + + add_check( + "full_episode_token_ids_present", + "pass" if token_ids else "fail", + f"token_count={len(token_ids)}", + ) + add_check( + "mask_length_matches_token_count", + "pass" if token_ids and len(mask) == len(token_ids) else "fail", + f"mask_count={len(mask)} token_count={len(token_ids)}", + ) + add_check( + "logprobs_length_matches_token_count", + "pass" if token_ids and len(logprobs) == len(token_ids) else "fail", + f"logprobs_count={len(logprobs)} token_count={len(token_ids)}", + ) + + if token_ids and len(logprobs) == len(token_ids): + generated_positions: set[int] = set() + for idx, trace in enumerate(token_turn_traces): + prompt_len = len(trace.get("prompt_ids") or []) + prefill_len = int((model_request_traces[idx] if idx < len(model_request_traces) else {}).get("assistant_prefill_len") or 0) + completion_logprobs = trace.get("completion_logprobs") or [] + for offset in range(len(completion_logprobs)): + generated_positions.add(prompt_len + prefill_len + offset) + + generated_present = all( + pos < len(logprobs) and logprobs[pos] is not None + for pos in generated_positions + ) + prompt_positions_empty = all( + logprobs[pos] is None + for pos in range(len(logprobs)) + if pos not in generated_positions + ) + add_check( + "generated_logprobs_present", + "pass" if generated_present else "fail", + f"generated_positions={sorted(generated_positions)}", + ) + add_check( + "prompt_positions_have_no_logprobs", + "pass" if prompt_positions_empty else "fail", + "only generated-token positions should carry logprobs", + ) + else: + add_check("generated_logprobs_present", "skip", "full_episode.logprobs unavailable") + add_check("prompt_positions_have_no_logprobs", "skip", "full_episode.logprobs unavailable") + + return { + "target": _validation_target(model_id=model_id, tokenizer_name=tokenizer_name, visual=visual), + "passed": not failing_checks, + "failing_checks": failing_checks, + "skipped_checks": skipped_checks, + "turn_count": len(token_turn_traces), + "prompt_lengths": prompt_lengths, + "completion_lengths": completion_lengths, + "image_counts": image_counts, + "token_count": len(token_ids), + "mask_count": len(mask), + "validation_checks": validation_checks, + } + + +def enrich_rows(results: list[EvaluationRow], tokenizer_name: str, *, model_id: str, visual: bool): """Enrich rows with detokenized tokens, full-episode view, status, and score.""" from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) + try: + tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True) + except Exception as exc: + logger.warning("Tokenizer load failed for %s: %s", tokenizer_name, exc) + tokenizer = None for row in results: extra = row.execution_metadata.extra @@ -125,8 +392,10 @@ def enrich_rows(results: list[EvaluationRow], tokenizer_name: str): prompt_ids = trace.get("prompt_ids") or [] completion_ids = trace.get("completion_ids") or [] all_ids = [int(x) for x in prompt_ids] + [int(x) for x in completion_ids] - detok_tokens = [_detok(tokenizer, tid) for tid in all_ids] - trace["detokenized_tokens"] = detok_tokens + if tokenizer is not None: + trace["detokenized_tokens"] = [_detok(tokenizer, tid) for tid in all_ids] + trace["prompt_detokenized_tokens"] = [_detok(tokenizer, tid) for tid in prompt_ids] + trace["completion_detokenized_tokens"] = [_detok(tokenizer, tid) for tid in completion_ids] trace["prompt_len"] = len(prompt_ids) trace["completion_len"] = len(completion_ids) @@ -145,22 +414,33 @@ def enrich_rows(results: list[EvaluationRow], tokenizer_name: str): logprobs_arr: list[float | None] = [None] * full_len for k in range(num_turns): turn_prompt_len = len(token_turn_traces[k].get("prompt_ids") or []) + prefill_len = int((mrt[k] if k < len(mrt) else {}).get("assistant_prefill_len") or 0) turn_lp = token_turn_traces[k].get("completion_logprobs") or [] for j, lp in enumerate(turn_lp): - pos = turn_prompt_len + j + pos = turn_prompt_len + prefill_len + j if pos < full_len: logprobs_arr[pos] = lp - full_detok = [_detok(tokenizer, tid) for tid in full_ids] - extra["full_episode"] = { "token_ids": full_ids, "mask": mask, "logprobs": logprobs_arr, - "detokenized_tokens": full_detok, "num_turns": num_turns, "first_prompt_len": first_prompt_len, } + if tokenizer is not None: + extra["full_episode"]["detokenized_tokens"] = [_detok(tokenizer, tid) for tid in full_ids] + + validation_summary = _build_row_validation_details( + row, + model_id=model_id, + tokenizer_name=tokenizer_name, + visual=visual, + ) + extra["validation_summary"] = { + key: value for key, value in validation_summary.items() if key != "validation_checks" + } + extra["validation_checks"] = validation_summary["validation_checks"] # Set rollout_status and evaluation_result has_error = bool(extra.get("rollout_error")) @@ -190,20 +470,368 @@ def write_to_default_logger(results: list[EvaluationRow]): logger.info("Wrote %d rows to default logger (sqlite)", len(results)) +def configure_default_logger(*, logs_db_path: str, reset_logs: bool = False): + """Point eval-protocol's lazy default logger at a specific sqlite DB.""" + from eval_protocol.dataset_logger.sqlite_dataset_logger_adapter import SqliteDatasetLoggerAdapter + + adapter = SqliteDatasetLoggerAdapter(db_path=logs_db_path) + if reset_logs: + adapter._store.delete_all_rows() + logger.info("Cleared existing eval-protocol rows from %s", logs_db_path) + + default_logger._logger = adapter + logger.info("Using eval-protocol sqlite db at %s", logs_db_path) + + +def _reconstruct_turn_image_urls(row: EvaluationRow) -> list[str]: + dataset_info = row.input_metadata.dataset_info if isinstance(row.input_metadata.dataset_info, dict) else {} + env_context = dataset_info.get("environment_context") if isinstance(dataset_info, dict) else {} + if not isinstance(env_context, dict): + return [] + + extra = row.execution_metadata.extra if isinstance(row.execution_metadata.extra, dict) else {} + tool_call_traces = extra.get("tool_call_traces") or [] + token_turn_traces = extra.get("token_turn_traces") or [] + if not token_turn_traces: + return [] + + env = build_frozen_lake_tool_env(environment_context=env_context, max_steps=max(1, len(tool_call_traces) + 1)) + env.reset() + images = [env.render_image_data_url()] + for trace in tool_call_traces[: max(0, len(token_turn_traces) - 1)]: + action = trace.get("action") + if not isinstance(action, str): + break + env.step(action) + images.append(env.render_image_data_url()) + return images[: len(token_turn_traces)] + + +def _format_token_piece(value: object) -> str: + if value is None: + return "" + text = str(value) + return text.encode("unicode_escape").decode("ascii") + + +def _mask_class(mask_value: int) -> str: + if mask_value <= 0: + return "tok tok-prompt" + palette_index = ((mask_value - 1) % 6) + 1 + return f"tok tok-turn-{palette_index}" + + +def _status_chip(status: str) -> str: + normalized = str(status or "skip").lower() + if normalized == "pass": + return "chip chip-pass" + if normalized == "fail": + return "chip chip-fail" + return "chip chip-skip" + + +def build_debug_report_html(results: list[EvaluationRow]) -> str: + sections: list[str] = [] + for row in results: + extra = row.execution_metadata.extra if isinstance(row.execution_metadata.extra, dict) else {} + full_episode = extra.get("full_episode") or {} + token_ids = [int(x) for x in (full_episode.get("token_ids") or [])] + mask = [int(x) for x in (full_episode.get("mask") or [])] + logprobs = list(full_episode.get("logprobs") or []) + detok = list(full_episode.get("detokenized_tokens") or []) + token_turn_traces = extra.get("token_turn_traces") or [] + tool_call_traces = extra.get("tool_call_traces") or [] + model_request_traces = extra.get("model_request_traces") or [] + validation_summary = extra.get("validation_summary") or {} + validation_checks = extra.get("validation_checks") or [] + image_urls = _reconstruct_turn_image_urls(row) + rollout_id = "" + if row.execution_metadata and row.execution_metadata.rollout_id: + rollout_id = str(row.execution_metadata.rollout_id) + score = row.evaluation_result.score if row.evaluation_result else None + + token_spans: list[str] = [] + token_rows: list[str] = [] + for idx, token_id in enumerate(token_ids): + token_text = detok[idx] if idx < len(detok) else "" + mask_value = mask[idx] if idx < len(mask) else 0 + logprob_value = logprobs[idx] if idx < len(logprobs) else None + tooltip = html.escape( + f"idx={idx} id={token_id} mask={mask_value} logprob={logprob_value} text={_format_token_piece(token_text)}" + ) + token_spans.append( + f'{html.escape(_format_token_piece(token_text) or "")}' + ) + token_rows.append( + html.escape( + f"{idx:05d} id={token_id:<8} mask={mask_value:<3} logprob={logprob_value!s:<12} text={_format_token_piece(token_text)}" + ) + ) + + turn_rows: list[str] = [] + for turn_index, trace in enumerate(token_turn_traces): + request_trace = model_request_traces[turn_index] if turn_index < len(model_request_traces) else {} + tool_trace = tool_call_traces[turn_index] if turn_index < len(tool_call_traces) else {} + image_html = "" + if turn_index < len(image_urls): + image_html = f'turn {turn_index + 1} input image' + turn_meta = html.escape( + json.dumps( + { + "image_count": request_trace.get("image_count"), + "assistant_turn_len": request_trace.get("assistant_turn_len"), + "tool_suffix_len": request_trace.get("tool_suffix_len"), + "prompt_token_count": request_trace.get("prompt_token_count"), + }, + indent=2, + ) + ) + turn_rows.append( + "\n".join( + [ + "
", + f"Turn {turn_index + 1}: prompt={len(trace.get('prompt_ids') or [])} completion={len(trace.get('completion_ids') or [])} action={html.escape(str(tool_trace.get('action') or ''))} reward={html.escape(str(tool_trace.get('reward') or ''))}", + image_html, + f"
{turn_meta}
", + f"
{html.escape(str(request_trace.get('prompt_text') or ''))}
", + f"
{html.escape(json.dumps(trace, indent=2))}
", + "
", + ] + ) + ) + + validation_rows = [ + ( + f'
  • ' + f'{html.escape(str(check.get("status") or "skip").upper())} ' + f'{html.escape(str(check.get("name") or ""))}: ' + f'{html.escape(str(check.get("detail") or ""))}
    • ' + ) + for check in validation_checks + ] + validation_passed = bool(validation_summary.get("passed")) + validation_badge = ( + 'PASS' + if validation_passed else 'FAIL' + ) + + sections.append( + "\n".join( + [ + "
    ", + f"

    {html.escape(str(row.input_metadata.row_id))}

    ", + f"

    rollout_id: {html.escape(rollout_id)} | mode: {html.escape(str(extra.get('observation_mode') or ''))} | score: {html.escape(str(score))}

    ", + f"

    turns: {len(token_turn_traces)} | tokens: {len(token_ids)} | error: {html.escape(str(extra.get('rollout_error') or ''))}

    ", + f"

    validation: {validation_badge} target: {html.escape(str(validation_summary.get('target') or ''))} image_counts: {html.escape(str(validation_summary.get('image_counts') or []))}

    ", + "

    Validation Checks

    ", + "
      " + "".join(validation_rows) + "
    ", + "

    Masked Token Stream

    ", + "
    " + "".join(token_spans) + "
    ", + "
    Token Table
    " + "\n".join(token_rows) + "
    ", + "

    Per-Turn Trace

    ", + "\n".join(turn_rows), + "
    ", + ] + ) + ) + + return "\n".join( + [ + "", + "", + "", + "", + "FrozenLake Visual Rollout Report", + "", + "", + "", + "

    FrozenLake Visual Rollout Report

    ", + "

    This report is generated from eval-protocol rows after mask enrichment. Prompt tokens are gray; model output tokens are color-coded by turn index.

    ", + "\n".join(sections), + "", + "", + ] + ) + + +def write_debug_report(results: list[EvaluationRow], report_path: str): + report_file = Path(report_path) + report_file.parent.mkdir(parents=True, exist_ok=True) + report_file.write_text(build_debug_report_html(results), encoding="utf-8") + logger.info("Wrote debug report to %s", report_file) + + +def _dashboard_urls(port: int) -> tuple[str, str]: + base = f"http://localhost:{port}" + return f"{base}/table", f"{base}/frozen-lake-report" + + +def _schedule_browser_open(port: int, *, report_path: str) -> None: + table_url, report_url = _dashboard_urls(port) + + def _open() -> None: + logger.info("Opening logs UI at %s", table_url) + try: + webbrowser.open(table_url) + except Exception as exc: + logger.warning("Failed to open browser for %s: %s", table_url, exc) + if Path(report_path).exists(): + logger.info("FrozenLake report available at %s", report_url) + + opener = threading.Timer(1.0, _open) + opener.daemon = True + opener.start() + + +def serve_logs_dashboard(port: int, *, logs_db_path: str, report_path: str, open_browser: bool = False): + """Start the eval-protocol dashboard against the local sqlite logger. + + We bypass ``serve_logs()`` because the installed ``ep`` path and the default + uvicorn lifespan startup can both hang in this local environment. + """ + import uvicorn + from fastapi.responses import HTMLResponse + from eval_protocol.utils.browser_utils import write_pid_file + from eval_protocol.utils.logs_server import LogsServer + + async def _run(): + configure_default_logger(logs_db_path=logs_db_path) + server = LogsServer(port=port) + + if Path(report_path).exists(): + @server.app.get("/frozen-lake-report") + async def frozen_lake_report(): + return HTMLResponse(Path(report_path).read_text(encoding="utf-8")) + + report_route = server.app.router.routes.pop() + server.app.router.routes.insert(0, report_route) + + try: + server.start_loops() + config = uvicorn.Config( + server.app, + host=server.host, + port=server.port, + log_level="info", + lifespan="off", + ) + uvicorn_server = uvicorn.Server(config) + write_pid_file(os.getpid(), server.port) + await uvicorn_server.serve() + finally: + server.evaluation_watcher.stop() + server.websocket_manager.stop_broadcast_loop() + + table_url, report_url = _dashboard_urls(port) + logger.info("Starting eval-protocol dashboard on port %d...", port) + logger.info("Logs table: %s", table_url) + logger.info("FrozenLake report: %s", report_url) + if open_browser: + _schedule_browser_open(port, report_path=report_path) + asyncio.run(_run()) + + def main(): parser = argparse.ArgumentParser() parser.add_argument("--num-rollouts", type=int, default=4) + parser.add_argument("--rows-per-seed", type=int, default=1) parser.add_argument("--port", type=int, default=8765) parser.add_argument("--no-serve", action="store_true") + parser.add_argument("--serve-only", action="store_true") + parser.add_argument("--reset-logs", action="store_true") + parser.add_argument("--logs-db", default=DEFAULT_LOGS_DB_PATH) + parser.add_argument("--report-path", default=DEFAULT_REPORT_PATH) + parser.add_argument( + "--model-id", + default=os.environ.get("FROZEN_LAKE_VERIFY_MODEL_ID", "accounts/fireworks/models/qwen3-8b"), + ) + parser.add_argument("--tokenizer-model", default=None) + parser.add_argument("--visual", action="store_true") + parser.add_argument("--allow-plaintext-action-fallback", action="store_true") + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--seed", action="append", type=int, default=[]) + parser.add_argument("--open-browser", action=argparse.BooleanOptionalAction, default=True) args = parser.parse_args() + resolved_tokenizer_model = _resolve_tokenizer_model( + model_id=args.model_id, + tokenizer_model=args.tokenizer_model, + ) + + if args.serve_only: + serve_logs_dashboard( + args.port, + logs_db_path=args.logs_db, + report_path=args.report_path, + open_browser=args.open_browser, + ) + return - results = asyncio.run(run_rollouts(num_rollouts=args.num_rollouts)) + configure_default_logger(logs_db_path=args.logs_db, reset_logs=args.reset_logs) + + results = asyncio.run( + run_rollouts( + num_rollouts=args.num_rollouts, + rows_per_seed=args.rows_per_seed, + model_id=args.model_id, + tokenizer=resolved_tokenizer_model, + visual=args.visual, + allow_plaintext_action_fallback=args.allow_plaintext_action_fallback, + temperature=args.temperature, + seed_values=list(args.seed or []), + ) + ) if not results: logger.error("No rollouts completed") return logger.info("Enriching rows...") - enrich_rows(results, "Qwen/Qwen3-8B") + enrich_rows(results, resolved_tokenizer_model, model_id=args.model_id, visual=args.visual) + write_debug_report(results, args.report_path) + + rewards = [] + for row in results: + extra = row.execution_metadata.extra or {} + step_rewards = extra.get("step_rewards") or [] + rewards.append(1.0 if step_rewards and float(step_rewards[-1]) > 0 else 0.0) + if rewards: + validation_passes = 0 + for row in results: + extra = row.execution_metadata.extra or {} + validation_summary = extra.get("validation_summary") or {} + if validation_summary.get("passed"): + validation_passes += 1 + logger.info( + "Rollout summary: n=%d success_rate=%.3f avg_reward=%.3f visual=%s model=%s tokenizer=%s temperature=%.3f validation_pass_rows=%d", + len(rewards), + sum(rewards) / len(rewards), + sum(rewards) / len(rewards), + args.visual, + args.model_id, + resolved_tokenizer_model, + float(args.temperature), + validation_passes, + ) write_to_default_logger(results) @@ -211,9 +839,12 @@ def main(): logger.info("Done (--no-serve). Data written to default logger.") return - logger.info("Starting eval-protocol dashboard on port %d...", args.port) - from eval_protocol.utils.logs_server import serve_logs - serve_logs(port=args.port) + serve_logs_dashboard( + args.port, + logs_db_path=args.logs_db, + report_path=args.report_path, + open_browser=args.open_browser, + ) if __name__ == "__main__": diff --git a/training/pyproject.toml b/training/pyproject.toml index dc12fbab..2e2cfba1 100644 --- a/training/pyproject.toml +++ b/training/pyproject.toml @@ -9,12 +9,14 @@ requires-python = ">=3.11" dependencies = [ "fireworks-ai>=1.0.0a36,<2", "tinker-cookbook>=0.1.0", - "eval-protocol>=0.3.21", + "eval-protocol>=0.3.23", "tqdm", "torch", "datasets", "tiktoken", "numpy", + "transformers", + "pillow", "python-dotenv", "wandb", ] @@ -22,6 +24,7 @@ dependencies = [ [project.optional-dependencies] dev = [ "pytest", + "pytest-timeout", "math-verify", ] @@ -31,3 +34,7 @@ include = ["training*"] [tool.pytest.ini_options] testpaths = ["training/tests"] +markers = [ + "e2e: tests that hit remote Fireworks infrastructure", + "timeout: per-test timeout budget", +] diff --git a/training/recipes/dpo_loop.py b/training/recipes/dpo_loop.py index 4695f65a..5d92ee84 100644 --- a/training/recipes/dpo_loop.py +++ b/training/recipes/dpo_loop.py @@ -44,7 +44,6 @@ ReconnectableClient, wandb_log, setup_wandb, - extract_text, setup_resume, wandb_finish, validate_config, @@ -53,7 +52,9 @@ setup_deployment, create_trainer_job, load_preference_dataset, - find_common_prefix_length, + build_renderer, + render_preference_pair, + resolve_renderer_name, ) from fireworks.training.sdk.deployment import DEFAULT_DELTA_COMPRESSION from fireworks.training.sdk.weight_syncer import WeightSyncer @@ -98,6 +99,7 @@ class Config: base_model: str = "accounts/fireworks/models/qwen3-8b" dataset: str = "" tokenizer_model: str = "" # HuggingFace model name for client-side tokenization + renderer_name: str = "" beta: float = 0.1 learning_rate: float = 1e-5 @@ -129,44 +131,35 @@ class Config: def _tokenize_pair( example: dict[str, Any], tokenizer: Any, + renderer: Any, max_seq_len: int, ) -> dict[str, Any] | None: """Tokenize a single preference pair. Returns None if invalid, 'filtered' if too long.""" - chosen_text = extract_text(example["chosen"]) - rejected_text = extract_text(example["rejected"]) - if not chosen_text or not rejected_text: + pair = render_preference_pair( + example["chosen"], + example["rejected"], + renderer=renderer, + tokenizer=tokenizer, + ) + if pair is None: return None - - chosen_tokens = tokenizer.encode(chosen_text) - rejected_tokens = tokenizer.encode(rejected_text) - if len(chosen_tokens) > max_seq_len or len(rejected_tokens) > max_seq_len: + if len(pair.chosen_tokens) > max_seq_len or len(pair.rejected_tokens) > max_seq_len: return "filtered" - if len(chosen_tokens) < 2 or len(rejected_tokens) < 2: - return None - prompt_len = find_common_prefix_length(chosen_tokens, rejected_tokens) return { - "chosen_tokens": chosen_tokens, - "rejected_tokens": rejected_tokens, - "prompt_len": prompt_len, + "chosen_tokens": pair.chosen_tokens, + "rejected_tokens": pair.rejected_tokens, + "response_start": pair.response_start, + "chosen_datum": pair.chosen_datum, + "rejected_datum": pair.rejected_datum, } -def _make_datum(tokens: list[int]) -> tinker.Datum: - return tinker.Datum( - model_input=tinker.ModelInput.from_ints(tokens[:-1]), - loss_fn_inputs={ - "target_tokens": tinker.TensorData( - data=tokens[1:], dtype="int64", shape=[len(tokens) - 1] - ) - }, - ) - - async def _cache_ref_logprobs( raw_data: list[dict[str, Any]], reference: ReconnectableClient, tokenizer: Any, + renderer: Any, max_seq_len: int, concurrency: int = 16, batch_size: int = 1, @@ -181,7 +174,7 @@ async def _cache_ref_logprobs( tokenized: list[tuple[int, dict[str, Any]]] = [] filtered_count = 0 for i, example in enumerate(raw_data): - result = _tokenize_pair(example, tokenizer, max_seq_len) + result = _tokenize_pair(example, tokenizer, renderer, max_seq_len) if result == "filtered": filtered_count += 1 elif result is not None: @@ -198,8 +191,8 @@ async def _process_batch( ) -> list[tuple[int, dict[str, Any]]]: datums: list[tinker.Datum] = [] for _, pair_data in batch: - datums.append(_make_datum(pair_data["chosen_tokens"])) - datums.append(_make_datum(pair_data["rejected_tokens"])) + datums.append(pair_data["chosen_datum"]) + datums.append(pair_data["rejected_datum"]) async with semaphore: fwd = await asyncio.to_thread( @@ -211,9 +204,11 @@ async def _process_batch( results.append((idx, { "chosen_tokens": pair_data["chosen_tokens"], "rejected_tokens": pair_data["rejected_tokens"], + "chosen_datum": pair_data["chosen_datum"], + "rejected_datum": pair_data["rejected_datum"], "ref_chosen": fwd.loss_fn_outputs[2 * j]["logprobs"].data, "ref_rejected": fwd.loss_fn_outputs[2 * j + 1]["logprobs"].data, - "prompt_len": pair_data["prompt_len"], + "response_start": pair_data["response_start"], })) return results @@ -250,12 +245,11 @@ def _flush_batch( response_starts: list[int] = [] for cached in batch_pairs: - response_start = max(0, cached["prompt_len"] - 1) - datums.append(_make_datum(cached["chosen_tokens"])) - datums.append(_make_datum(cached["rejected_tokens"])) + datums.append(cached["chosen_datum"]) + datums.append(cached["rejected_datum"]) ref_chosen_list.append(cached["ref_chosen"]) ref_rejected_list.append(cached["ref_rejected"]) - response_starts.append(response_start) + response_starts.append(cached["response_start"]) loss_fn = make_batch_dpo_loss_fn( ref_chosen_list, ref_rejected_list, response_starts, beta, @@ -423,6 +417,14 @@ def _signal_handler(signum, frame): if cfg.deployment.deployment_id: setup_deployment(deploy_mgr, cfg.deployment, cfg.base_model, cfg.infra) + profile = None + if cfg.infra.training_shape_id: + profile = rlor_mgr.resolve_training_profile(cfg.infra.training_shape_id) + + if profile and cfg.max_seq_len is None: + cfg.max_seq_len = profile.max_supported_context_length + logger.info("max_seq_len from training shape: %d", cfg.max_seq_len) + if cfg.max_seq_len is None: raise ValueError( "max_seq_len is required. Set it in Config, or use a training shape " @@ -445,6 +447,7 @@ def _signal_handler(signum, frame): rlor_mgr, base_model=cfg.base_model, infra=cfg.infra, + profile=profile, lora_rank=cfg.lora_rank, max_seq_len=cfg.max_seq_len, learning_rate=cfg.learning_rate, @@ -456,6 +459,7 @@ def _signal_handler(signum, frame): rlor_mgr, base_model=cfg.base_model, infra=cfg.infra, + profile=profile, lora_rank=cfg.lora_rank, max_seq_len=cfg.max_seq_len, learning_rate=cfg.learning_rate, @@ -489,6 +493,11 @@ def _signal_handler(signum, frame): import transformers tokenizer = transformers.AutoTokenizer.from_pretrained(cfg.tokenizer_model, trust_remote_code=True) + renderer = build_renderer(tokenizer, cfg.tokenizer_model, cfg.renderer_name) + logger.info( + "Using renderer=%s for preference tokenization", + resolve_renderer_name(cfg.tokenizer_model, cfg.renderer_name), + ) raw_data = load_preference_dataset(cfg.dataset, cfg.max_pairs) if not raw_data: @@ -497,7 +506,7 @@ def _signal_handler(signum, frame): logger.info("Computing reference logprobs for %d pairs...", len(raw_data)) ref_cache, filtered_count = asyncio.run( _cache_ref_logprobs( - raw_data, reference, tokenizer, cfg.max_seq_len, + raw_data, reference, tokenizer, renderer, cfg.max_seq_len, concurrency=cfg.ref_cache_concurrency, batch_size=cfg.ref_cache_batch_size, ) diff --git a/training/recipes/orpo_loop.py b/training/recipes/orpo_loop.py index 24320d44..f2a11533 100644 --- a/training/recipes/orpo_loop.py +++ b/training/recipes/orpo_loop.py @@ -59,7 +59,9 @@ make_orpo_loss_fn, create_trainer_job, load_preference_dataset, - find_common_prefix_length, + build_renderer, + render_preference_pair, + resolve_renderer_name, ) logger = logging.getLogger(__name__) @@ -74,6 +76,7 @@ class Config: base_model: str = "accounts/fireworks/models/qwen3-235b-a22b-instruct-2507" dataset: str = "" tokenizer_model: str = "" + renderer_name: str = "" orpo_lambda: float = 1.0 learning_rate: float = 1e-5 @@ -143,6 +146,14 @@ def _signal_handler(signum, frame): api_key=api_key, account_id=account, base_url=base_url ) + profile = None + if cfg.infra.training_shape_id: + profile = rlor_mgr.resolve_training_profile(cfg.infra.training_shape_id) + + if profile and cfg.max_seq_len is None: + cfg.max_seq_len = profile.max_supported_context_length + logger.info("max_seq_len from training shape: %d", cfg.max_seq_len) + if cfg.max_seq_len is None: raise ValueError( "max_seq_len is required. Set it in Config, or use a training shape " @@ -153,6 +164,7 @@ def _signal_handler(signum, frame): rlor_mgr, base_model=cfg.base_model, infra=cfg.infra, + profile=profile, lora_rank=cfg.lora_rank, max_seq_len=cfg.max_seq_len, learning_rate=cfg.learning_rate, @@ -173,6 +185,11 @@ def _signal_handler(signum, frame): tokenizer = transformers.AutoTokenizer.from_pretrained( cfg.tokenizer_model, trust_remote_code=True ) + renderer = build_renderer(tokenizer, cfg.tokenizer_model, cfg.renderer_name) + logger.info( + "Using renderer=%s for preference tokenization", + resolve_renderer_name(cfg.tokenizer_model, cfg.renderer_name), + ) raw_data = load_preference_dataset(cfg.dataset, cfg.max_pairs) if not raw_data: @@ -183,33 +200,28 @@ def _signal_handler(signum, frame): filtered_count = 0 for example in raw_data: - chosen_msgs = example["chosen"].get("messages", []) - rejected_msgs = example["rejected"].get("messages", []) - if not chosen_msgs or not rejected_msgs: - continue - - chosen_tokens = tokenizer.apply_chat_template( - chosen_msgs, tokenize=True, return_dict=False - ) - rejected_tokens = tokenizer.apply_chat_template( - rejected_msgs, tokenize=True, return_dict=False + pair = render_preference_pair( + example["chosen"], + example["rejected"], + renderer=renderer, + tokenizer=tokenizer, ) + if pair is None: + continue if ( - len(chosen_tokens) > cfg.max_seq_len - or len(rejected_tokens) > cfg.max_seq_len + len(pair.chosen_tokens) > cfg.max_seq_len + or len(pair.rejected_tokens) > cfg.max_seq_len ): filtered_count += 1 continue - if len(chosen_tokens) < 2 or len(rejected_tokens) < 2: - continue - - prompt_len = find_common_prefix_length(chosen_tokens, rejected_tokens) pair_cache.append( { - "chosen_tokens": chosen_tokens, - "rejected_tokens": rejected_tokens, - "prompt_len": prompt_len, + "chosen_tokens": pair.chosen_tokens, + "rejected_tokens": pair.rejected_tokens, + "response_start": pair.response_start, + "chosen_datum": pair.chosen_datum, + "rejected_datum": pair.rejected_datum, } ) @@ -246,32 +258,9 @@ def _signal_handler(signum, frame): for pair in pair_cache: chosen_tokens = pair["chosen_tokens"] rejected_tokens = pair["rejected_tokens"] - response_start = max(0, pair["prompt_len"] - 1) - - chosen_datum = tinker.Datum( - model_input=tinker.ModelInput.from_ints(chosen_tokens[:-1]), - loss_fn_inputs={ - "target_tokens": tinker.TensorData( - data=chosen_tokens[1:], - dtype="int64", - shape=[len(chosen_tokens) - 1], - ) - }, - ) - rejected_datum = tinker.Datum( - model_input=tinker.ModelInput.from_ints(rejected_tokens[:-1]), - loss_fn_inputs={ - "target_tokens": tinker.TensorData( - data=rejected_tokens[1:], - dtype="int64", - shape=[len(rejected_tokens) - 1], - ) - }, - ) - - loss_fn = make_orpo_loss_fn(response_start, cfg.orpo_lambda) + loss_fn = make_orpo_loss_fn(pair["response_start"], cfg.orpo_lambda) result = client.forward_backward_custom( - [chosen_datum, rejected_datum], loss_fn + [pair["chosen_datum"], pair["rejected_datum"]], loss_fn ).result() metrics = result.metrics diff --git a/training/recipes/rl_loop.py b/training/recipes/rl_loop.py index 825fa647..792c1b02 100644 --- a/training/recipes/rl_loop.py +++ b/training/recipes/rl_loop.py @@ -47,6 +47,7 @@ compute_advantages, create_trainer_job, load_jsonl_dataset, + prepare_sampling_messages, ) from fireworks.training.sdk.deployment import DeploymentSampler from training.utils.rl import PromptGroup @@ -57,7 +58,7 @@ from training.utils.rl.dapo import DAPOConfig from training.utils.rl.gspo import GSPOConfig from training.utils.rl.cispo import CISPOConfig -from training.utils.rl.train import TrainStepFns, run_rl_loop +from training.utils.rl.train import MinibatchTrainFns, run_rl_loop from training.utils.rl.losses import build_loss_fn, combine_prompt_groups from training.utils.rl.metrics import compute_step_metrics from training.utils.rl.router_replay import build_r3_routing_matrices @@ -382,7 +383,7 @@ def _signal_handler(signum, frame): async def sample_one_prompt(row: dict) -> PromptGroup | None: """Sample completions for one prompt and return a PromptGroup.""" messages = row.get("messages", []) - input_messages = [m for m in messages if m.get("role") != "assistant"] + input_messages = prepare_sampling_messages(messages) if not input_messages: return None @@ -485,30 +486,35 @@ def ref_forward(groups: list[PromptGroup]) -> None: ] idx += n - def train_step( - step: int, - prompt_groups: list[PromptGroup], - loop_stats: dict | None = None, - ) -> tuple[int, dict]: - """ref_forward + fwd_bwd + optim_step + hotload + metrics (1:1).""" + def fwd_bwd_one(prompt_groups: list[PromptGroup]): + """One minibatch forward/backward call after reference forward.""" + if not prompt_groups: + raise ValueError("fwd_bwd_one requires at least one prompt group") import time as _t - t0 = _t.time() - ref_forward(prompt_groups) - logger.info("[step %d] ref_forward: done (%.1fs)", step + 1, _t.time() - t0) - data, adv, ref_lp, prompt_lens, inf_lp = combine_prompt_groups(prompt_groups) t0 = _t.time() prox_fwd = policy.forward(data, "cross_entropy") prox_lp = [prox_fwd.loss_fn_outputs[i]["logprobs"].data for i in range(len(data))] - logger.info("[step %d] prox_forward: done (%.1fs)", step + 1, _t.time() - t0) + logger.info("prox_forward: done (%.1fs)", _t.time() - t0) t0 = _t.time() fwd_bwd_result = policy.forward_backward_custom( data, loss_builder(adv, ref_lp, prompt_lens, inf_lp, prox_lp), ) - logger.info("[step %d] fwd_bwd: done (%.1fs)", step + 1, _t.time() - t0) + logger.info("fwd_bwd: done (%.1fs)", _t.time() - t0) + return fwd_bwd_result + + def finish_step( + step: int, + prompt_groups: list[PromptGroup], + fwd_bwd_results: list, + n_accum: int, + loop_stats: dict | None = None, + ) -> tuple[int, dict]: + """optim_step + hotload + metrics after all minibatches in a step.""" + import time as _t t0 = _t.time() optim_result = policy.optim_step(adam_params) @@ -530,9 +536,9 @@ def train_step( metrics = compute_step_metrics( prompt_groups=prompt_groups, - fwd_bwd_results=[fwd_bwd_result], + fwd_bwd_results=fwd_bwd_results, optim_result=optim_result, - n_accum=1, + n_accum=n_accum, timing_metrics=flush_timing(), loop_stats=loop_stats, completions_per_prompt=completions_per_prompt, @@ -560,18 +566,23 @@ def _loop_metrics_callback(loop_metrics: dict) -> None: """Called by run_rl_loop after each train step with loop-level metrics.""" wandb_log(loop_metrics, step=loop_metrics.get("train/step", 0)) - train_fns = TrainStepFns(train_step=train_step) + train_fns = MinibatchTrainFns( + ref_forward_batch=ref_forward, + fwd_bwd_one=fwd_bwd_one, + finish_step=finish_step, + ) all_rows = dataset * cfg.epochs global_step = asyncio.run(run_rl_loop( sample_fns=(sample_one_prompt(row) for row in all_rows), - train_fns=train_fns, + minibatch_fns=train_fns, prompt_groups_per_step=prompt_groups_per_step, max_concurrent=cfg.max_concurrent, dynamic_filter_fn=should_accept, global_step=step_offset, metrics_callback=_loop_metrics_callback, + completions_per_prompt=completions_per_prompt, )) # -- Final checkpoint ---------------------------------------------------- diff --git a/training/recipes/sft_loop.py b/training/recipes/sft_loop.py index 78c57eae..b04e7bc8 100644 --- a/training/recipes/sft_loop.py +++ b/training/recipes/sft_loop.py @@ -42,7 +42,11 @@ validate_config, log_metrics_json, create_trainer_job, - make_batch_sft_loss_fn, + make_batch_weighted_sft_loss_fn, + build_renderer, + parse_train_on_what, + render_messages_to_datum, + resolve_renderer_name, ) from training.utils.timer import timer, flush_timing @@ -60,6 +64,8 @@ class Config: base_model: str = "accounts/fireworks/models/qwen3-8b" dataset: str = "" tokenizer_model: str = "" # HuggingFace model name for chat template, e.g. "Qwen/Qwen3-1.7B" + renderer_name: str = "" + train_on_what: str = "all_assistant_messages" learning_rate: float = 1e-4 epochs: int = 3 @@ -150,6 +156,13 @@ def _signal_handler(signum, frame): # -- Prepare data ------------------------------------------------------ try: tokenizer = transformers.AutoTokenizer.from_pretrained(cfg.tokenizer_model, trust_remote_code=True) + renderer = build_renderer(tokenizer, cfg.tokenizer_model, cfg.renderer_name) + train_on_what = parse_train_on_what(cfg.train_on_what) + logger.info( + "Using renderer=%s train_on_what=%s", + resolve_renderer_name(cfg.tokenizer_model, cfg.renderer_name), + train_on_what.value, + ) raw_data: List[Dict[str, Any]] = [] with open(cfg.dataset) as f: @@ -162,26 +175,23 @@ def _signal_handler(signum, frame): break logger.info("Loaded %d examples from %s", len(raw_data), cfg.dataset) - training_data: List[Dict[str, Any]] = [] + training_data: List[tinker.Datum] = [] filtered_count = 0 for row in raw_data: messages = row.get("messages", []) if not messages: continue - full_tokens = tokenizer.apply_chat_template(messages, tokenize=True, return_dict=False) - if len(full_tokens) > cfg.max_seq_len or len(full_tokens) < 2: + rendered = render_messages_to_datum( + messages, + renderer=renderer, + train_on_what=train_on_what, + ) + if len(rendered.token_ids) > cfg.max_seq_len or len(rendered.token_ids) < 2: filtered_count += 1 continue - prompt_messages = [m for m in messages if m.get("role") != "assistant"] - prompt_tokens = tokenizer.apply_chat_template( - prompt_messages, - tokenize=True, - add_generation_prompt=True, - return_dict=False, - ) - training_data.append({"tokens": full_tokens, "prompt_len": len(prompt_tokens)}) + training_data.append(rendered.datum) if filtered_count > 0: logger.info( @@ -206,24 +216,13 @@ def _signal_handler(signum, frame): agg_loss_sum = 0.0 agg_resp_tokens = 0 - def _create_datum(tokens: list[int]) -> tinker.Datum: - return tinker.Datum( - model_input=tinker.ModelInput.from_ints(tokens[:-1]), - loss_fn_inputs={ - "target_tokens": tinker.TensorData(data=tokens[1:], dtype="int64", shape=[len(tokens) - 1]), - }, - ) - - def _flush_batch(batch_buf: list[dict], step: int, accum: int) -> tuple[int, int]: + def _flush_batch(batch_buf: list[tinker.Datum], step: int, accum: int) -> tuple[int, int]: """Send a batch through forward_backward_custom and return (step, accum).""" nonlocal agg_loss_sum, agg_resp_tokens - datums = [_create_datum(ex["tokens"]) for ex in batch_buf] - prompt_counts = [ex["prompt_len"] for ex in batch_buf] - - loss_fn = make_batch_sft_loss_fn(prompt_counts) + loss_fn = make_batch_weighted_sft_loss_fn() with timer("fwd_bwd"): - result = client.forward_backward_custom(datums, loss_fn) + result = client.forward_backward_custom(batch_buf, loss_fn) fwd_metrics = result.metrics agg_loss_sum += fwd_metrics.get("ce_loss_sum", 0.0) @@ -271,7 +270,7 @@ def _flush_batch(batch_buf: list[dict], step: int, accum: int) -> tuple[int, int return step, accum for _epoch in range(cfg.epochs): - batch_buffer: list[dict] = [] + batch_buffer: list[tinker.Datum] = [] for ex in training_data: batch_buffer.append(ex) if len(batch_buffer) >= batch_size: diff --git a/training/tests/smoke_test/conftest.py b/training/tests/smoke_test/conftest.py new file mode 100644 index 00000000..4f343f16 --- /dev/null +++ b/training/tests/smoke_test/conftest.py @@ -0,0 +1,75 @@ +"""Shared fixtures for remote smoke tests on small Qwen3-4B shapes.""" + +from __future__ import annotations + +import os + +import pytest + +from fireworks.training.sdk.deployment import DeploymentManager +from fireworks.training.sdk.trainer import TrainerJobManager + +DEFAULT_SMOKE_BASE_MODEL = "accounts/fireworks/models/qwen3-4b" +DEFAULT_SMOKE_TOKENIZER_MODEL = "Qwen/Qwen3-4B" +DEFAULT_SMOKE_TRAINING_SHAPE = "ts-qwen3-4b-smoke-v1" +DEFAULT_SMOKE_BASE_URL = "https://dev.api.fireworks.ai" + + +def _get_env(name: str, default: str | None = None) -> str | None: + return os.environ.get(name, default) + + +@pytest.fixture(scope="session") +def smoke_base_model() -> str: + return _get_env("FIREWORKS_SMOKE_BASE_MODEL", DEFAULT_SMOKE_BASE_MODEL) + + +@pytest.fixture(scope="session") +def smoke_tokenizer_model() -> str: + return _get_env("FIREWORKS_SMOKE_TOKENIZER_MODEL", DEFAULT_SMOKE_TOKENIZER_MODEL) + + +@pytest.fixture(scope="session") +def smoke_training_shape() -> str: + return _get_env("FIREWORKS_SMOKE_TRAINING_SHAPE", DEFAULT_SMOKE_TRAINING_SHAPE) + + +@pytest.fixture(scope="session") +def smoke_custom_image_tag() -> str | None: + return _get_env("FIREWORKS_CUSTOM_IMAGE_TAG") + + +@pytest.fixture(scope="session") +def smoke_sdk_managers(): + api_key = _get_env("FIREWORKS_API_KEY") + if not api_key: + pytest.skip("FIREWORKS_API_KEY not set") + + account_id = _get_env("FIREWORKS_ACCOUNT_ID") + if not account_id: + pytest.skip("FIREWORKS_ACCOUNT_ID not set") + + base_url = _get_env("FIREWORKS_BASE_URL", DEFAULT_SMOKE_BASE_URL) + inference_url = _get_env("FIREWORKS_INFERENCE_URL", base_url) + hotload_api_url = _get_env("FIREWORKS_HOTLOAD_API_URL", base_url) + + additional_headers = {} + gateway_secret = _get_env("FIREWORKS_GATEWAY_SECRET") + if gateway_secret: + additional_headers["X-Fireworks-Gateway-Secret"] = gateway_secret + + rlor_mgr = TrainerJobManager( + api_key=api_key, + account_id=account_id, + base_url=base_url, + additional_headers=additional_headers or None, + ) + deploy_mgr = DeploymentManager( + api_key=api_key, + account_id=account_id, + base_url=base_url, + inference_url=inference_url, + hotload_api_url=hotload_api_url, + additional_headers=additional_headers or None, + ) + return rlor_mgr, deploy_mgr diff --git a/training/tests/smoke_test/test_dpo_smoke.py b/training/tests/smoke_test/test_dpo_smoke.py new file mode 100644 index 00000000..83278551 --- /dev/null +++ b/training/tests/smoke_test/test_dpo_smoke.py @@ -0,0 +1,75 @@ +"""Minimal DPO smoke test on Qwen3-4B.""" + +from __future__ import annotations + +import json +import os +import tempfile + +import pytest + + +def _make_preference_dataset(path: str, num_pairs: int = 4) -> None: + with open(path, "w") as f: + for i in range(num_pairs): + row = { + "chosen": { + "messages": [ + {"role": "user", "content": f"What is {i} plus {i}?"}, + {"role": "assistant", "content": f"The answer is {i + i}."}, + ] + }, + "rejected": { + "messages": [ + {"role": "user", "content": f"What is {i} plus {i}?"}, + {"role": "assistant", "content": f"I think it is {i * 3}."}, + ] + }, + } + f.write(json.dumps(row) + "\n") + + +@pytest.mark.e2e +@pytest.mark.timeout(3600) +def test_dpo_smoke( + smoke_sdk_managers, + smoke_base_model, + smoke_tokenizer_model, + smoke_training_shape, + smoke_custom_image_tag, +): + from training.recipes.dpo_loop import Config, main + from training.utils import DeployConfig, HotloadConfig, InfraConfig, WandBConfig + + rlor_mgr, deploy_mgr = smoke_sdk_managers + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + dataset_path = f.name + + try: + _make_preference_dataset(dataset_path, num_pairs=4) + config = Config( + base_model=smoke_base_model, + dataset=dataset_path, + tokenizer_model=smoke_tokenizer_model, + beta=0.1, + learning_rate=1e-5, + epochs=1, + batch_size=1, + grad_accum=1, + max_pairs=4, + infra=InfraConfig( + training_shape_id=smoke_training_shape, + skip_validations=True, + custom_image_tag=smoke_custom_image_tag, + ), + deployment=DeployConfig(), + hotload=HotloadConfig(hot_load_interval=0, dcp_save_interval=0), + wandb=WandBConfig(), + ) + + metrics = main(config, rlor_mgr=rlor_mgr, deploy_mgr=deploy_mgr) + assert isinstance(metrics, dict) + assert metrics["steps"] >= 2, f"Expected >= 2 steps, got {metrics['steps']}" + finally: + os.unlink(dataset_path) diff --git a/training/tests/smoke_test/test_grpo_smoke.py b/training/tests/smoke_test/test_grpo_smoke.py index b611567b..0bdfba40 100644 --- a/training/tests/smoke_test/test_grpo_smoke.py +++ b/training/tests/smoke_test/test_grpo_smoke.py @@ -11,7 +11,7 @@ FIREWORKS_BASE_URL (optional, defaults to https://dev.api.fireworks.ai) Usage: - pytest training/tests/e2e/smoke_test/test_grpo_smoke.py -v -s + pytest training/tests/smoke_test/test_grpo_smoke.py -v -s """ from __future__ import annotations @@ -23,20 +23,12 @@ import pytest -from fireworks.training.sdk.trainer import TrainerJobManager -from fireworks.training.sdk.deployment import DeploymentManager - logging.basicConfig( level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s", force=True, ) -BASE_MODEL = "accounts/fireworks/models/qwen3-4b" -TOKENIZER_MODEL = "Qwen/Qwen3-4B" -TRAINING_SHAPE = "ts-qwen3-4b-smoke-v1" -DEFAULT_BASE_URL = "https://dev.api.fireworks.ai" - def _make_prompt_dataset(path: str, n: int = 4) -> None: with open(path, "w") as f: @@ -54,50 +46,21 @@ def _constant_reward(_completion: str, _row: dict) -> float: return 1.0 -def _get_sdk_managers(): - api_key = os.environ.get("FIREWORKS_API_KEY") - if not api_key: - pytest.skip("FIREWORKS_API_KEY not set") - - account_id = os.environ.get("FIREWORKS_ACCOUNT_ID") - if not account_id: - pytest.skip("FIREWORKS_ACCOUNT_ID not set") - - base_url = os.environ.get("FIREWORKS_BASE_URL", DEFAULT_BASE_URL) - inference_url = os.environ.get("FIREWORKS_INFERENCE_URL", base_url) - hotload_api_url = os.environ.get("FIREWORKS_HOTLOAD_API_URL", base_url) - - additional_headers = {} - gateway_secret = os.environ.get("FIREWORKS_GATEWAY_SECRET") - if gateway_secret: - additional_headers["X-Fireworks-Gateway-Secret"] = gateway_secret - - rlor_mgr = TrainerJobManager( - api_key=api_key, - account_id=account_id, - base_url=base_url, - additional_headers=additional_headers or None, - ) - deploy_mgr = DeploymentManager( - api_key=api_key, - account_id=account_id, - base_url=base_url, - inference_url=inference_url, - hotload_api_url=hotload_api_url, - additional_headers=additional_headers or None, - ) - return rlor_mgr, deploy_mgr - - @pytest.mark.e2e @pytest.mark.timeout(3600) -def test_grpo_smoke(): +def test_grpo_smoke( + smoke_sdk_managers, + smoke_base_model, + smoke_tokenizer_model, + smoke_training_shape, + smoke_custom_image_tag, +): """2-step GRPO on Qwen3-4B: train, hotload, train again, cleanup.""" from training.utils import InfraConfig, DeployConfig, HotloadConfig, WandBConfig from training.recipes.rl_loop import Config, main import training.recipes.rl_loop as rl_mod - rlor_mgr, deploy_mgr = _get_sdk_managers() + rlor_mgr, deploy_mgr = smoke_sdk_managers rl_mod.reward_fn = _constant_reward @@ -108,7 +71,7 @@ def test_grpo_smoke(): _make_prompt_dataset(dataset_path, n=4) config = Config( - base_model=BASE_MODEL, + base_model=smoke_base_model, dataset=dataset_path, learning_rate=1e-4, kl_beta=0, @@ -118,11 +81,12 @@ def test_grpo_smoke(): max_rows=2, epochs=1, infra=InfraConfig( - training_shape_id=TRAINING_SHAPE, + training_shape_id=smoke_training_shape, skip_validations=True, + custom_image_tag=smoke_custom_image_tag, ), deployment=DeployConfig( - tokenizer_model=TOKENIZER_MODEL, + tokenizer_model=smoke_tokenizer_model, ), hotload=HotloadConfig( hot_load_interval=1, diff --git a/training/tests/smoke_test/test_sft_smoke.py b/training/tests/smoke_test/test_sft_smoke.py new file mode 100644 index 00000000..509f7d70 --- /dev/null +++ b/training/tests/smoke_test/test_sft_smoke.py @@ -0,0 +1,64 @@ +"""Minimal SFT smoke test on Qwen3-4B.""" + +from __future__ import annotations + +import json +import os +import tempfile + +import pytest + + +def _make_chat_dataset(path: str, num_examples: int = 4) -> None: + with open(path, "w") as f: + for i in range(num_examples): + row = { + "messages": [ + {"role": "user", "content": f"What is {i} plus {i}?"}, + {"role": "assistant", "content": f"The answer is {i + i}."}, + ] + } + f.write(json.dumps(row) + "\n") + + +@pytest.mark.e2e +@pytest.mark.timeout(3600) +def test_sft_smoke( + smoke_sdk_managers, + smoke_base_model, + smoke_tokenizer_model, + smoke_training_shape, + smoke_custom_image_tag, +): + from training.recipes.sft_loop import Config, main + from training.utils import InfraConfig, WandBConfig + + rlor_mgr, _deploy_mgr = smoke_sdk_managers + + with tempfile.NamedTemporaryFile(mode="w", suffix=".jsonl", delete=False) as f: + dataset_path = f.name + + try: + _make_chat_dataset(dataset_path, num_examples=4) + config = Config( + base_model=smoke_base_model, + dataset=dataset_path, + tokenizer_model=smoke_tokenizer_model, + learning_rate=1e-4, + epochs=1, + batch_size=2, + grad_accum=1, + max_examples=4, + infra=InfraConfig( + training_shape_id=smoke_training_shape, + skip_validations=True, + custom_image_tag=smoke_custom_image_tag, + ), + wandb=WandBConfig(), + ) + + metrics = main(config, rlor_mgr=rlor_mgr) + assert isinstance(metrics, dict) + assert metrics["steps"] >= 2, f"Expected >= 2 steps, got {metrics['steps']}" + finally: + os.unlink(dataset_path) diff --git a/training/tests/test_smoke_imports.py b/training/tests/test_smoke_imports.py index e260c391..523b26d6 100644 --- a/training/tests/test_smoke_imports.py +++ b/training/tests/test_smoke_imports.py @@ -68,7 +68,7 @@ def test_util_imports(module: str): # ── Example modules ───────────────────────────────────────────────────────── EXAMPLE_MODULES_WITH_ENV = [ - ("training.examples.deepmath.train_deepmath", "math_verify", {"FIREWORKS_API_KEY": "test"}), + ("training.examples.deepmath_rl.train_deepmath", "math_verify", {"FIREWORKS_API_KEY": "test"}), ] @@ -127,9 +127,6 @@ def test_tinker_types(attr: str): ("fireworks.training.sdk.trainer", "TrainerJobManager"), ("fireworks.training.sdk.trainer", "TrainerServiceEndpoint"), ("fireworks.training.sdk.trainer", "TrainingShapeProfile"), - ("fireworks.training.sdk.errors", "DOCS_HOTLOAD"), - ("fireworks.training.sdk.errors", "DOCS_API_KEYS"), - ("fireworks.training.sdk.errors", "DOCS_DEPLOYMENTS"), ("fireworks.training.sdk.errors", "format_sdk_error"), ("fireworks.training.sdk.errors", "request_with_retries"), ] diff --git a/training/tests/unit/test_frozen_lake_verify_validation.py b/training/tests/unit/test_frozen_lake_verify_validation.py new file mode 100644 index 00000000..eaaf9ae6 --- /dev/null +++ b/training/tests/unit/test_frozen_lake_verify_validation.py @@ -0,0 +1,193 @@ +from __future__ import annotations + +import json + +from eval_protocol.models import EvaluationRow, InputMetadata, Message + +from training.examples.frozen_lake.frozen_lake_schema import FROZEN_LAKE_TOOLS +from training.examples.frozen_lake.verify_rollout import build_debug_report_html, enrich_rows + + +def _build_visual_row(*, row_id: str, image_counts: list[int]) -> EvaluationRow: + row = EvaluationRow( + input_metadata=InputMetadata( + row_id=row_id, + dataset_info={"environment_context": {"desc": ["SF", "FG"]}}, + ) + ) + row.tools = list(FROZEN_LAKE_TOOLS) + row.messages = [ + Message( + role="system", + content="You are an RL policy for FrozenLake.\nAlways respond with exactly one tool call, no text.", + ), + Message( + role="user", + content="You are playing FrozenLake. The image shows the current grid and the text below gives the same observation.\n[S] F\n F G\n\nReply with exactly one token: LEFT, DOWN, RIGHT, or UP.", + ), + Message( + role="assistant", + content="RIGHT", + tool_calls=[ + { + "id": "call_1", + "type": "function", + "function": {"name": "lake_move", "arguments": '{"action":"RIGHT"}'}, + } + ], + ), + Message( + role="tool", + name="lake_move", + tool_call_id="call_1", + content=json.dumps( + { + "observation": " S [F]\n F G", + "action": "RIGHT", + "reward": 0.0, + "terminated": False, + "truncated": False, + }, + separators=(",", ":"), + ), + ), + Message( + role="assistant", + content="DOWN", + tool_calls=[ + { + "id": "call_2", + "type": "function", + "function": {"name": "lake_move", "arguments": '{"action":"DOWN"}'}, + } + ], + ), + Message( + role="tool", + name="lake_move", + tool_call_id="call_2", + content=json.dumps( + { + "observation": " S F\n F [G]", + "action": "DOWN", + "reward": 1.0, + "terminated": True, + "truncated": False, + }, + separators=(",", ":"), + ), + ), + ] + row.execution_metadata.extra = { + "observation_mode": "image", + "step_rewards": [0.0, 1.0], + "token_turn_traces": [ + { + "step_index": 1, + "prompt_ids": [10, 11, 12], + "completion_ids": [99, 101, 90], + "completion_logprobs": [-0.1], + "prompt_text": "prompt:initial", + }, + { + "step_index": 2, + "prompt_ids": [10, 11, 12, 99, 101, 90, 91, 92], + "completion_ids": [99, 102, 90], + "completion_logprobs": [-0.2], + "prompt_text": "prompt:stitched", + }, + ], + "model_request_traces": [ + { + "step_index": 1, + "prompt_ids": [10, 11, 12, 99], + "prompt_token_count": 4, + "prompt_text": "prompt:initial", + "assistant_prefill_len": 1, + "assistant_turn_len": 3, + "tool_suffix_len": 3, + "image_count": image_counts[0], + }, + { + "step_index": 2, + "prompt_ids": [10, 11, 12, 99, 101, 90, 91, 92, 99], + "prompt_token_count": 9, + "prompt_text": "prompt:stitched", + "assistant_prefill_len": 1, + "assistant_turn_len": 3, + "tool_suffix_len": 0, + "image_count": image_counts[1], + }, + ], + "tool_call_traces": [ + {"step_index": 1, "tool_name": "lake_move", "action": "RIGHT", "reward": 0.0}, + {"step_index": 2, "tool_name": "lake_move", "action": "DOWN", "reward": 1.0}, + ], + } + return row + + +def test_enrich_rows_adds_passing_visual_validation_summary(): + row = _build_visual_row(row_id="passing", image_counts=[1, 2]) + + enrich_rows( + [row], + "missing-tokenizer-for-test", + model_id="accounts/fireworks/models/kimi-k2p5-vl", + visual=True, + ) + + extra = row.execution_metadata.extra or {} + validation_summary = extra["validation_summary"] + validation_checks = extra["validation_checks"] + full_episode = extra["full_episode"] + + assert validation_summary["target"] == "kimi-k2.5-vl" + assert validation_summary["passed"] is True + assert validation_summary["failing_checks"] == [] + assert validation_summary["image_counts"] == [1, 2] + assert len(full_episode["token_ids"]) == len(full_episode["mask"]) == len(full_episode["logprobs"]) + assert full_episode["mask"][3] == 1 + assert full_episode["logprobs"][3] is None + assert full_episode["logprobs"][4] == -0.1 + assert full_episode["mask"][5] == 1 + assert full_episode["logprobs"][5] is None + assert any(check["name"] == "image_counts_increment_per_turn" and check["status"] == "pass" for check in validation_checks) + assert any( + check["name"] == "kimi_instant_prompt_contains_empty_think_block" and check["status"] == "pass" + for check in validation_checks + ) + + +def test_enrich_rows_flags_broken_image_count_validation(): + row = _build_visual_row(row_id="failing", image_counts=[1, 1]) + + enrich_rows( + [row], + "missing-tokenizer-for-test", + model_id="accounts/fireworks/models/kimi-k2p5-vl", + visual=True, + ) + + extra = row.execution_metadata.extra or {} + validation_summary = extra["validation_summary"] + + assert validation_summary["passed"] is False + assert "image_counts_increment_per_turn" in validation_summary["failing_checks"] + + +def test_debug_report_renders_validation_section(): + row = _build_visual_row(row_id="report_row", image_counts=[1, 2]) + + enrich_rows( + [row], + "missing-tokenizer-for-test", + model_id="accounts/fireworks/models/kimi-k2p5-vl", + visual=True, + ) + + report_html = build_debug_report_html([row]) + + assert "Validation Checks" in report_html + assert "kimi-k2.5-vl" in report_html + assert "image_counts_increment_per_turn" in report_html diff --git a/training/tests/unit/test_frozen_lake_visual_rollout.py b/training/tests/unit/test_frozen_lake_visual_rollout.py new file mode 100644 index 00000000..b8a0809d --- /dev/null +++ b/training/tests/unit/test_frozen_lake_visual_rollout.py @@ -0,0 +1,320 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from eval_protocol.models import EvaluationRow, InputMetadata +from eval_protocol.pytest.types import RolloutProcessorConfig + +from training.examples.frozen_lake.frozen_lake_rollout import ( + FireworksV1ImageCompletionsClient, + FrozenLakeToolRolloutProcessor, + _build_visual_user_prompt, + build_frozen_lake_tool_call_parser, +) +from training.examples.frozen_lake.frozen_lake_schema import FROZEN_LAKE_TOOLS + + +class _FakeImageClient: + requested_prompt_ids: list[list[int]] = [] + requested_image_counts: list[int] = [] + requested_prompt_texts: list[str] = [] + + def __init__(self, **_: object) -> None: + self._call_count = 0 + + def build_prompt_token_ids(self, *, messages, tools): + assert tools, "visual rollout should still pass tool specs into prompt construction" + assert messages[-1]["role"] == "user" + assert isinstance(messages[-1]["content"], list) + return [10, 11, 12] + + def decode_token_ids(self, *, token_ids): + normalized = list(token_ids) + if normalized == [10, 11, 12]: + return "prompt:initial" + if normalized == [91, 92]: + return "suffix:tool" + return f"prompt:{normalized}" + + def build_prompt_text(self, *, messages, tools): + assert tools, "visual rollout should still pass tool specs into prompt construction" + assert messages[-1]["role"] == "user" + return "prompt:initial" + + def encode_assistant_turn_suffix(self): + return [90] + + def assistant_turn_suffix_text(self): + return "<|im_end|>" + + def build_tool_response_suffix_token_ids(self, *, tool_message): + assert tool_message["role"] == "tool" + assert isinstance(tool_message["content"], list) + return [91, 92] + + def build_tool_response_suffix_text(self, *, tool_message): + assert tool_message["role"] == "tool" + return "suffix:tool" + + async def create_completion_from_prompt_ids(self, *, prompt_token_ids, prompt_text=None, images, tools): + del tools + self._call_count += 1 + type(self).requested_prompt_ids.append(list(prompt_token_ids)) + type(self).requested_image_counts.append(len(images)) + type(self).requested_prompt_texts.append(str(prompt_text or "")) + + action = "RIGHT" if self._call_count == 1 else "DOWN" + completion_id = 100 + self._call_count + tool_call_id = f"tc_{self._call_count}" + return { + "choices": [ + { + "message": { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "lake_move", + "arguments": f'{{"action":"{action}"}}', + }, + } + ], + }, + "finish_reason": "stop", + "raw_output": {}, + } + ], + "usage": { + "prompt_tokens": len(prompt_token_ids), + "completion_tokens": 1, + "total_tokens": len(prompt_token_ids) + 1, + }, + "prompt_ids": list(prompt_token_ids), + "prompt_text": f"prompt:{prompt_token_ids}", + "completion_ids": [completion_id], + "completion_text": f"action:{action}", + "completion_logprobs": [-0.1], + "finish_reason": "stop", + "raw_output": {}, + } + + async def close(self) -> None: + return None + + +class _CapturingImageClient(_FakeImageClient): + init_kwargs: list[dict[str, object]] = [] + + def __init__(self, **kwargs: object) -> None: + type(self).init_kwargs.append(dict(kwargs)) + super().__init__(**kwargs) + + +def test_visual_user_prompt_includes_text_observation_without_placeholder(): + rendered = _build_visual_user_prompt( + prompt_template="Inspect the image and choose one action.", + observation="Current state:\nSFFF", + default_prompt="unused", + ) + + assert "Inspect the image and choose one action." in rendered + assert "Current textual observation:" in rendered + assert "Current state:\nSFFF" in rendered + + +def test_frozen_lake_tool_call_parser_accepts_json_schema_output_with_suffix(): + parser = build_frozen_lake_tool_call_parser(model_id="accounts/fireworks/models/kimi-k2p5") + + parsed = parser( + '{\n"tool_calls":[{"name":"lake_move","arguments":{"action":"RIGHT"}}]\n}<|im_end|>', + [1, 2, 3], + None, + ) + + assert parsed["parser"] == "json_schema" + assert parsed["assistant_content"] == "" + assert parsed["parsed_tool_call"].name == "lake_move" + assert parsed["parsed_tool_call"].arguments == {"action": "RIGHT"} + + +def test_frozen_lake_tool_call_parser_accepts_kimi_native_output(): + parser = build_frozen_lake_tool_call_parser(model_id="accounts/fireworks/models/kimi-k2p5") + + parsed = parser( + '<|tool_call_begin|>functions.lake_move:0<|tool_call_argument_begin|>{"action":"DOWN"}<|tool_call_end|><|tool_calls_section_end|>', + [1, 2, 3], + None, + ) + + assert parsed["parsed_tool_call"].name == "lake_move" + assert parsed["parsed_tool_call"].arguments == {"action": "DOWN"} + + +def test_kimi_visual_prompt_includes_empty_think_block_when_thinking_disabled(): + client = FireworksV1ImageCompletionsClient( + model_id="accounts/fireworks/models/kimi-k2p5", + tokenizer_name_or_path="moonshotai/Kimi-K2.5", + enable_thinking=False, + ) + try: + messages = [ + {"role": "system", "content": "You are an RL policy for FrozenLake."}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, + {"type": "text", "text": "obs"}, + ], + }, + ] + prompt_text = client.build_prompt_text(messages=messages, tools=list(FROZEN_LAKE_TOOLS)) + prompt_ids = client.build_prompt_token_ids(messages=messages, tools=list(FROZEN_LAKE_TOOLS)) + decoded_prompt = client.decode_token_ids(token_ids=prompt_ids) + + assert prompt_text.endswith("") + assert "" in decoded_prompt + assert "" in decoded_prompt + assert decoded_prompt.rstrip().endswith("    ") + finally: + asyncio.run(client.close()) + + +def test_kimi_visual_tool_suffix_includes_empty_think_block_when_thinking_disabled(): + client = FireworksV1ImageCompletionsClient( + model_id="accounts/fireworks/models/kimi-k2p5", + tokenizer_name_or_path="moonshotai/Kimi-K2.5", + enable_thinking=False, + ) + try: + tool_message = { + "role": "tool", + "name": "lake_move", + "tool_call_id": "functions.lake_move:0", + "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}], + } + suffix_text = client.build_tool_response_suffix_text(tool_message=tool_message) + suffix_ids = client.build_tool_response_suffix_token_ids(tool_message=tool_message) + decoded_suffix = client.decode_token_ids(token_ids=suffix_ids) + + assert suffix_text.endswith("") + assert "" in decoded_suffix + assert "" in decoded_suffix + assert decoded_suffix.rstrip().endswith("    ") + finally: + asyncio.run(client.close()) + + +@pytest.mark.asyncio +async def test_visual_rollout_preserves_prior_completion_tokens(monkeypatch): + _FakeImageClient.requested_prompt_ids = [] + _FakeImageClient.requested_image_counts = [] + _FakeImageClient.requested_prompt_texts = [] + monkeypatch.setattr( + "training.examples.frozen_lake.frozen_lake_rollout.FireworksV1ImageCompletionsClient", + _FakeImageClient, + ) + + processor = FrozenLakeToolRolloutProcessor( + model_id="accounts/fireworks/models/kimi-k2p5", + tokenizer_name_or_path="moonshotai/Kimi-K2.5", + observation_mode="image", + logprobs=True, + ) + row = EvaluationRow( + input_metadata=InputMetadata( + row_id="visual_seed", + dataset_info={"environment_context": {"desc": ["SF", "FG"]}}, + ) + ) + config = RolloutProcessorConfig( + completion_params={"model": "accounts/fireworks/models/kimi-k2p5"}, + mcp_config_path="", + steps=4, + semaphore=asyncio.Semaphore(1), + logger=None, + ) + + [task] = processor([row], config) + result = await task + + extra = result.execution_metadata.extra or {} + token_turn_traces = extra["token_turn_traces"] + + assert len(token_turn_traces) == 2 + assert token_turn_traces[0]["prompt_ids"] == [10, 11, 12] + assert token_turn_traces[0]["completion_ids"] == [163595, 101, 90] + assert token_turn_traces[1]["prompt_ids"] == [10, 11, 12, 163595, 101, 90, 91, 92] + assert token_turn_traces[1]["completion_ids"] == [163595, 102, 90] + + assert _FakeImageClient.requested_prompt_ids == [ + [10, 11, 12, 163595], + [10, 11, 12, 163595, 101, 90, 91, 92, 163595], + ] + assert _FakeImageClient.requested_prompt_texts == [ + "prompt:initial<|tool_calls_section_begin|>", + "prompt:initial<|tool_calls_section_begin|>action:RIGHT<|im_end|>suffix:tool<|tool_calls_section_begin|>", + ] + assert _FakeImageClient.requested_image_counts == [1, 2] + assert extra["model_request_traces"][0]["assistant_turn_len"] == 3 + assert extra["model_request_traces"][1]["assistant_turn_len"] == 3 + assert extra["observation_mode"] == "image" + assert extra["tool_call_generation_mode"] == "prompt_only" + assert [message.role for message in result.messages] == [ + "system", + "user", + "assistant", + "tool", + "assistant", + "tool", + ] + assert str(result.messages[2].content).startswith("action:RIGHT") + assert isinstance(result.messages[1].content, list) + assert result.messages[1].content[0]["type"] == "image_url" + assert isinstance(result.messages[3].content, list) + assert result.messages[3].content[0]["type"] == "image_url" + + +@pytest.mark.asyncio +async def test_kimi_visual_rollout_keeps_raw_prompt_only_generation(monkeypatch): + _CapturingImageClient.init_kwargs = [] + monkeypatch.setattr( + "training.examples.frozen_lake.frozen_lake_rollout.FireworksV1ImageCompletionsClient", + _CapturingImageClient, + ) + + processor = FrozenLakeToolRolloutProcessor( + model_id="accounts/fireworks/models/kimi-k2p5", + tokenizer_name_or_path="moonshotai/Kimi-K2.5", + observation_mode="image", + logprobs=True, + ) + row = EvaluationRow( + input_metadata=InputMetadata( + row_id="visual_seed_json_schema", + dataset_info={"environment_context": {"desc": ["SF", "FG"]}}, + ) + ) + config = RolloutProcessorConfig( + completion_params={"model": "accounts/fireworks/models/kimi-k2p5"}, + mcp_config_path="", + steps=2, + semaphore=asyncio.Semaphore(1), + logger=None, + ) + + [task] = processor([row], config) + result = await task + + init_kwargs = _CapturingImageClient.init_kwargs[0] + request_params = init_kwargs["request_params"] + assert isinstance(request_params, dict) + assert request_params["thinking"] == {"type": "disabled"} + assert "response_format" not in request_params + assert "logit_bias" not in request_params + assert init_kwargs["max_tokens"] == 256 + assert (result.execution_metadata.extra or {})["tool_call_generation_mode"] == "prompt_only" diff --git a/training/tests/unit/test_supervised_rendering.py b/training/tests/unit/test_supervised_rendering.py new file mode 100644 index 00000000..a426aa10 --- /dev/null +++ b/training/tests/unit/test_supervised_rendering.py @@ -0,0 +1,163 @@ +from __future__ import annotations + +import pytest +import torch +from tinker_cookbook.renderers import TrainOnWhat + +from training.utils.losses import make_batch_weighted_sft_loss_fn +from training.utils.data import prepare_sampling_messages +from training.utils.supervised import ( + build_datum_from_token_mask, + render_preference_pair, + normalize_messages, + render_messages_to_datum, +) + + +class StubRenderer: + def __init__(self, tokens: list[int], weights: list[float]): + self.tokens = torch.tensor(tokens, dtype=torch.int64) + self.weights = torch.tensor(weights, dtype=torch.float32) + self.calls: list[tuple[list[dict], TrainOnWhat]] = [] + + def build_supervised_example(self, messages, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE): + self.calls.append((messages, train_on_what)) + return self.tokens, self.weights + + +class SequenceRenderer: + def __init__(self, outputs: list[tuple[list[int], list[float]]]): + self.outputs = [ + (torch.tensor(tokens, dtype=torch.int64), torch.tensor(weights, dtype=torch.float32)) + for tokens, weights in outputs + ] + self.calls: list[tuple[list[dict], TrainOnWhat]] = [] + + def build_supervised_example(self, messages, train_on_what=TrainOnWhat.LAST_ASSISTANT_MESSAGE): + self.calls.append((messages, train_on_what)) + return self.outputs[len(self.calls) - 1] + + +def test_render_messages_to_datum_preserves_multi_turn_weights(): + renderer = StubRenderer( + tokens=[10, 11, 12, 13, 14, 15, 16, 17, 18], + weights=[0, 0, 1, 1, 0, 0, 1, 1, 1], + ) + + rendered = render_messages_to_datum( + [ + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ], + renderer=renderer, + train_on_what="all_assistant_messages", + ) + + normalized_messages, train_on_what = renderer.calls[0] + assert [m["role"] for m in normalized_messages] == ["user", "assistant", "user", "assistant"] + assert train_on_what == TrainOnWhat.ALL_ASSISTANT_MESSAGES + + assert rendered.token_ids == [10, 11, 12, 13, 14, 15, 16, 17, 18] + assert rendered.token_weights == [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0] + assert rendered.datum.loss_fn_inputs["target_tokens"].data == [11, 12, 13, 14, 15, 16, 17, 18] + assert rendered.datum.loss_fn_inputs["weights"].data == [0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0] + + +def test_build_datum_from_token_mask_reuses_ui_mask_semantics(): + rendered = build_datum_from_token_mask( + token_ids=[100, 101, 102, 103, 104, 105], + token_mask=[0, 0, 1, 1, 0, 2], + include_loss_mask=True, + ) + + assert rendered.token_weights == [0.0, 0.0, 1.0, 1.0, 0.0, 1.0] + assert rendered.datum.loss_fn_inputs["target_tokens"].data == [101, 102, 103, 104, 105] + assert rendered.datum.loss_fn_inputs["weights"].data == [0.0, 1.0, 1.0, 0.0, 1.0] + assert rendered.datum.loss_fn_inputs["loss_mask"].data == [0.0, 1.0, 1.0, 0.0, 1.0] + + +def test_normalize_messages_supports_openai_tool_call_shape(): + normalized = normalize_messages( + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": { + "name": "lake_move", + "arguments": '{"action":"RIGHT"}', + }, + } + ], + } + ] + ) + + assert normalized[0]["tool_calls"] == [{"name": "lake_move", "args": {"action": "RIGHT"}}] + + +def test_weighted_sft_loss_uses_sparse_weights(): + datum_a = build_datum_from_token_mask( + token_ids=[10, 11, 12, 13], + token_mask=[0, 0, 1, 0], + ).datum + datum_b = build_datum_from_token_mask( + token_ids=[20, 21, 22], + token_mask=[0, 1, 1], + ).datum + + loss_fn = make_batch_weighted_sft_loss_fn() + loss, metrics = loss_fn( + [datum_a, datum_b], + [ + torch.tensor([-0.1, -0.2, -0.3], dtype=torch.float32), + torch.tensor([-0.4, -0.5], dtype=torch.float32), + ], + ) + + assert loss.item() == pytest.approx(1.1 / 3.0) + assert metrics["ce_loss_sum"] == pytest.approx(1.1) + assert metrics["response_tokens"] == pytest.approx(3.0) + assert metrics["weighted_tokens"] == pytest.approx(3.0) + + +def test_render_preference_pair_uses_shared_renderer_path(): + renderer = SequenceRenderer( + outputs=[ + ([1, 2, 3, 4, 5, 6], [0, 0, 0, 1, 1, 1]), + ([1, 2, 3, 9, 10], [0, 0, 0, 1, 1]), + ] + ) + + pair = render_preference_pair( + {"messages": [{"role": "user", "content": "u"}, {"role": "assistant", "content": "good"}]}, + {"messages": [{"role": "user", "content": "u"}, {"role": "assistant", "content": "bad"}]}, + renderer=renderer, + tokenizer=None, + ) + + assert pair is not None + assert pair.chosen_tokens == [1, 2, 3, 4, 5, 6] + assert pair.rejected_tokens == [1, 2, 3, 9, 10] + assert pair.response_start == 3 + assert pair.chosen_datum.loss_fn_inputs["target_tokens"].data == [2, 3, 4, 5, 6] + assert len(renderer.calls) == 2 + + +def test_prepare_sampling_messages_only_strips_trailing_assistant(): + prepared = prepare_sampling_messages( + [ + {"role": "system", "content": "sys"}, + {"role": "user", "content": "u1"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "u2"}, + {"role": "assistant", "content": "a2"}, + ] + ) + + assert [m["role"] for m in prepared] == ["system", "user", "assistant", "user"] diff --git a/training/tests/unit/test_validate_preflight.py b/training/tests/unit/test_validate_preflight.py index 8df357ed..e02a4059 100644 --- a/training/tests/unit/test_validate_preflight.py +++ b/training/tests/unit/test_validate_preflight.py @@ -6,6 +6,7 @@ import pytest +import training.utils.validation as validation_module from training.utils.validation import validate_preflight @@ -118,3 +119,24 @@ def test_passes(self): hot_load_interval=0, ) validate_preflight(args, fw_api_key="key", fw_account_id="acct") + + +def test_validation_falls_back_when_sdk_docs_constants_are_missing(monkeypatch): + class FakeErrors: + @staticmethod + def format_sdk_error(what, cause, solution, docs_url=None): + lines = [f"ERROR: {what}", f"Cause: {cause}", f"Solution: {solution}"] + if docs_url: + lines.append(f"Docs: {docs_url}") + return " | ".join(lines) + + monkeypatch.setattr(validation_module, "sdk_errors", FakeErrors()) + + assert validation_module._get_sdk_docs_url("DOCS_HOTLOAD", "fallback-hotload") == "fallback-hotload" + assert validation_module._get_sdk_docs_url("DOCS_API_KEYS", "fallback-keys") == "fallback-keys" + assert "fallback-keys" in validation_module._format_sdk_error( + "Missing FIREWORKS_API_KEY", + "No API key found.", + "Set FIREWORKS_API_KEY.", + docs_url=validation_module._get_sdk_docs_url("DOCS_API_KEYS", "fallback-keys"), + ) diff --git a/training/utils/__init__.py b/training/utils/__init__.py index 63cb1159..a59595d4 100644 --- a/training/utils/__init__.py +++ b/training/utils/__init__.py @@ -28,7 +28,20 @@ "make_orpo_loss_fn", "make_batch_dpo_loss_fn", "make_batch_sft_loss_fn", + "make_batch_weighted_sft_loss_fn", "make_sft_loss_fn", + "RenderedSupervisedDatum", + "RenderedPreferencePair", + "build_next_token_datum", + "build_datum_from_token_mask", + "build_datum_from_tokens_and_weights", + "build_renderer", + "normalize_messages", + "parse_train_on_what", + "render_preference_pair", + "render_messages_to_datum", + "resolve_renderer_name", + "prepare_sampling_messages", "setup_deployment", "setup_resume", "setup_training_client", @@ -49,6 +62,7 @@ load_jsonl_dataset, load_preference_dataset, find_common_prefix_length, + prepare_sampling_messages, ) from training.utils.infra import ( setup_deployment, @@ -73,6 +87,20 @@ make_orpo_loss_fn, make_batch_dpo_loss_fn, make_batch_sft_loss_fn, + make_batch_weighted_sft_loss_fn, +) +from training.utils.supervised import ( + RenderedPreferencePair, + RenderedSupervisedDatum, + build_datum_from_token_mask, + build_next_token_datum, + build_datum_from_tokens_and_weights, + build_renderer, + normalize_messages, + parse_train_on_what, + render_preference_pair, + render_messages_to_datum, + resolve_renderer_name, ) from training.utils.resume import setup_resume from training.utils.logging import ( diff --git a/training/utils/data.py b/training/utils/data.py index eaf4d498..0a468189 100644 --- a/training/utils/data.py +++ b/training/utils/data.py @@ -97,6 +97,19 @@ def extract_text(item: dict[str, Any]) -> str: return "" +def prepare_sampling_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """Preserve multi-turn chat history, dropping only trailing assistant turns. + + RL prompt datasets should contain the full prior conversation. If the row + accidentally includes a final assistant completion, strip only that tail so + sampling resumes from the latest non-assistant turn. + """ + prepared = list(messages) + while prepared and prepared[-1].get("role") == "assistant": + prepared.pop() + return prepared + + def find_common_prefix_length(tokens1: List[int], tokens2: List[int]) -> int: """Find the length of the longest common prefix between two token lists.""" min_len = min(len(tokens1), len(tokens2)) diff --git a/training/utils/losses.py b/training/utils/losses.py index ef03cf79..942120b2 100644 --- a/training/utils/losses.py +++ b/training/utils/losses.py @@ -245,3 +245,59 @@ def loss_fn( } return loss_fn + + +def make_batch_weighted_sft_loss_fn( +) -> Callable[[List[tinker.Datum], List[torch.Tensor]], Tuple[torch.Tensor, Dict[str, float]]]: + """Cross-entropy loss over per-token supervised weights stored on each datum. + + This is the renderer-safe path for multi-turn SFT. Each datum must include + ``loss_fn_inputs["weights"]`` aligned with ``target_tokens``. + """ + + def loss_fn( + data: List[tinker.Datum], + logprobs_list: List[torch.Tensor], + ) -> Tuple[torch.Tensor, Dict[str, float]]: + assert len(data) == len(logprobs_list) + + total_loss = torch.tensor(0.0) + total_weight = 0.0 + total_nll = 0.0 + + for datum, logprobs in zip(data, logprobs_list, strict=True): + weights_td = datum.loss_fn_inputs.get("weights") + if weights_td is None: + raise ValueError("Weighted SFT expects each datum to include loss_fn_inputs['weights'].") + + weights = weights_td.to_torch().to(device=logprobs.device, dtype=logprobs.dtype) + if len(weights) != len(logprobs): + raise ValueError( + f"weights/logprobs length mismatch: {len(weights)} != {len(logprobs)}" + ) + + sample_weight = float(weights.sum().item()) + if sample_weight <= 0: + continue + + sample_nll = -(logprobs * weights).sum() + total_loss = total_loss + sample_nll + total_weight += sample_weight + with torch.no_grad(): + total_nll += float(sample_nll.item()) + + avg_loss = total_loss / total_weight if total_weight > 0 else total_loss + with torch.no_grad(): + avg_nll = total_nll / total_weight if total_weight > 0 else 0.0 + ppl = torch.exp(torch.tensor(avg_nll)).item() + + return avg_loss, { + "ce_loss": avg_nll, + "ce_loss_sum": total_nll, + "ppl": ppl, + "response_tokens": total_weight, + "weighted_tokens": total_weight, + "batch_size": len(logprobs_list), + } + + return loss_fn diff --git a/training/utils/supervised.py b/training/utils/supervised.py new file mode 100644 index 00000000..3e404550 --- /dev/null +++ b/training/utils/supervised.py @@ -0,0 +1,317 @@ +"""Shared supervised-rendering helpers. + +This module gives cookbook training code one token-level representation for +supervised data: + +- text/tool conversations rendered via a Tinker renderer +- eval-protocol-style token trajectories with a per-token mask + +Both paths end in the same ``tinker.Datum`` schema with ``target_tokens`` and +token-level ``weights`` so training uses the same spans that the UI shows. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Any, Iterable, Mapping, Sequence + +import torch +import tinker +from tinker_cookbook.model_info import get_recommended_renderer_name +from tinker_cookbook.renderers import Message, Renderer, TrainOnWhat, get_renderer +from tinker_cookbook.supervised.common import datum_from_tokens_weights + + +@dataclass(frozen=True) +class RenderedSupervisedDatum: + """Rendered tokens, token-level weights, and the final training datum.""" + + token_ids: list[int] + token_weights: list[float] + datum: tinker.Datum + + +@dataclass(frozen=True) +class RenderedPreferencePair: + """Rendered chosen/rejected preference pair with a shared response boundary.""" + + chosen_tokens: list[int] + rejected_tokens: list[int] + response_start: int + chosen_datum: tinker.Datum + rejected_datum: tinker.Datum + + +def parse_train_on_what(value: str | TrainOnWhat) -> TrainOnWhat: + """Normalize ``train_on_what`` config into Tinker's enum.""" + if isinstance(value, TrainOnWhat): + return value + return TrainOnWhat(value.lower()) + + +def resolve_renderer_name( + tokenizer_model: str, + renderer_name: str = "", +) -> str: + """Choose the renderer used for message -> token rendering.""" + if renderer_name: + return renderer_name + try: + return get_recommended_renderer_name(tokenizer_model) + except Exception as exc: # pragma: no cover - message only + raise ValueError( + f"Could not infer a renderer for tokenizer_model={tokenizer_model!r}. " + "Set Config.renderer_name explicitly." + ) from exc + + +def build_renderer( + tokenizer: Any, + tokenizer_model: str, + renderer_name: str = "", +) -> Renderer: + """Construct the Tinker renderer used for supervised formatting.""" + return get_renderer(resolve_renderer_name(tokenizer_model, renderer_name), tokenizer) + + +def _normalize_tool_calls(tool_calls: Any) -> list[dict[str, Any]]: + """Normalize common tool-call shapes into Tinker's ``{\"name\", \"args\"}`` form.""" + normalized: list[dict[str, Any]] = [] + for tool_call in tool_calls or []: + if not isinstance(tool_call, Mapping): + raise TypeError(f"Unsupported tool call type: {type(tool_call)!r}") + + if isinstance(tool_call.get("name"), str) and isinstance(tool_call.get("args"), Mapping): + normalized.append({ + "name": tool_call["name"], + "args": dict(tool_call["args"]), + }) + continue + + function = tool_call.get("function") + if isinstance(function, Mapping) and isinstance(function.get("name"), str): + raw_args = function.get("arguments", {}) + if isinstance(raw_args, str): + parsed_args = json.loads(raw_args) if raw_args else {} + elif isinstance(raw_args, Mapping): + parsed_args = dict(raw_args) + else: + raise TypeError(f"Unsupported tool call arguments type: {type(raw_args)!r}") + normalized.append({ + "name": function["name"], + "args": parsed_args, + }) + continue + + raise ValueError(f"Unsupported tool call shape: {tool_call}") + return normalized + + +def _normalize_content(content: Any) -> str: + """Convert OpenAI-style message content into plain text for text renderers.""" + if content is None: + return "" + if isinstance(content, str): + return content + if isinstance(content, Mapping): + if isinstance(content.get("text"), str): + return content["text"] + raise TypeError(f"Unsupported message content mapping: {content}") + if isinstance(content, Sequence): + text_parts: list[str] = [] + for part in content: + if isinstance(part, str): + text_parts.append(part) + continue + if not isinstance(part, Mapping): + raise TypeError(f"Unsupported message content part: {part!r}") + part_type = part.get("type") + if part_type == "text" and isinstance(part.get("text"), str): + text_parts.append(part["text"]) + continue + raise ValueError( + "Multimodal content is not supported by the shared text renderer. " + "Use the token_ids + mask path for eval-protocol visual trajectories." + ) + return "".join(text_parts) + raise TypeError(f"Unsupported message content type: {type(content)!r}") + + +def normalize_messages(messages: Iterable[Mapping[str, Any]]) -> list[Message]: + """Normalize cookbook/eval-style messages into Tinker's message schema.""" + normalized: list[Message] = [] + for message in messages: + role = message.get("role") + if not isinstance(role, str): + raise ValueError(f"Message is missing a string role: {message}") + + normalized_message: Message = { + "role": role, + "content": _normalize_content(message.get("content")), + } + + tool_calls = message.get("tool_calls") + if tool_calls is not None: + normalized_message["tool_calls"] = _normalize_tool_calls(tool_calls) + + thinking = message.get("thinking") + if thinking is not None: + if not isinstance(thinking, str): + raise TypeError(f"Unsupported thinking value type: {type(thinking)!r}") + normalized_message["thinking"] = thinking + + trainable = message.get("trainable") + if trainable is not None: + normalized_message["trainable"] = bool(trainable) + + normalized.append(normalized_message) + + return normalized + + +def build_datum_from_tokens_and_weights( + token_ids: Sequence[int], + token_weights: Sequence[float], + *, + max_seq_len: int | None = None, + include_loss_mask: bool = False, +) -> RenderedSupervisedDatum: + """Build a weighted ``tinker.Datum`` from full tokens and per-token weights.""" + tokens = [int(x) for x in token_ids] + weights = [float(x) for x in token_weights] + if len(tokens) != len(weights): + raise ValueError(f"tokens/weights length mismatch: {len(tokens)} != {len(weights)}") + if len(tokens) < 2: + raise ValueError("Need at least 2 tokens to build a supervised datum.") + + if max_seq_len is not None: + tokens = tokens[:max_seq_len] + weights = weights[:max_seq_len] + if len(tokens) < 2: + raise ValueError("Truncation left fewer than 2 tokens.") + + token_tensor = torch.tensor(tokens, dtype=torch.int64) + weight_tensor = torch.tensor(weights, dtype=torch.float32) + datum = datum_from_tokens_weights(token_tensor, weight_tensor) + + if include_loss_mask: + shifted_weights = [float(x) for x in weights[1:]] + datum.loss_fn_inputs["loss_mask"] = tinker.TensorData( + data=shifted_weights, + dtype="float32", + shape=[len(shifted_weights)], + ) + + return RenderedSupervisedDatum( + token_ids=tokens, + token_weights=weights, + datum=datum, + ) + + +def build_next_token_datum(token_ids: Sequence[int]) -> tinker.Datum: + """Build a standard next-token prediction datum without token weights.""" + tokens = [int(x) for x in token_ids] + if len(tokens) < 2: + raise ValueError("Need at least 2 tokens to build a next-token datum.") + return tinker.Datum( + model_input=tinker.ModelInput.from_ints(tokens[:-1]), + loss_fn_inputs={ + "target_tokens": tinker.TensorData( + data=tokens[1:], + dtype="int64", + shape=[len(tokens) - 1], + ) + }, + ) + + +def build_datum_from_token_mask( + token_ids: Sequence[int], + token_mask: Sequence[int | float], + *, + max_seq_len: int | None = None, + include_loss_mask: bool = False, +) -> RenderedSupervisedDatum: + """Build a weighted datum from an eval-protocol-style per-token mask.""" + token_weights = [1.0 if float(mask_value) > 0 else 0.0 for mask_value in token_mask] + return build_datum_from_tokens_and_weights( + token_ids, + token_weights, + max_seq_len=max_seq_len, + include_loss_mask=include_loss_mask, + ) + + +def render_messages_to_datum( + messages: Sequence[Mapping[str, Any]], + *, + renderer: Renderer, + train_on_what: str | TrainOnWhat = TrainOnWhat.ALL_ASSISTANT_MESSAGES, + max_seq_len: int | None = None, + include_loss_mask: bool = False, +) -> RenderedSupervisedDatum: + """Render a multi-turn conversation into the shared weighted datum format.""" + normalized_messages = normalize_messages(messages) + tokens, weights = renderer.build_supervised_example( + normalized_messages, + train_on_what=parse_train_on_what(train_on_what), + ) + return build_datum_from_tokens_and_weights( + tokens.tolist(), + weights.tolist(), + max_seq_len=max_seq_len, + include_loss_mask=include_loss_mask, + ) + + +def _common_prefix_length(tokens_a: Sequence[int], tokens_b: Sequence[int]) -> int: + min_len = min(len(tokens_a), len(tokens_b)) + for idx in range(min_len): + if int(tokens_a[idx]) != int(tokens_b[idx]): + return idx + return min_len + + +def _render_preference_item_tokens( + item: Mapping[str, Any], + *, + renderer: Renderer, + tokenizer: Any, +) -> list[int]: + if "messages" in item: + messages = item.get("messages") or [] + if not messages: + return [] + return render_messages_to_datum(messages, renderer=renderer).token_ids + if isinstance(item.get("text"), str): + return [int(x) for x in tokenizer.encode(item["text"])] + return [] + + +def render_preference_pair( + chosen: Mapping[str, Any], + rejected: Mapping[str, Any], + *, + renderer: Renderer, + tokenizer: Any, + max_seq_len: int | None = None, +) -> RenderedPreferencePair | None: + """Render a chosen/rejected pair through the shared tokenizer path.""" + chosen_tokens = _render_preference_item_tokens(chosen, renderer=renderer, tokenizer=tokenizer) + rejected_tokens = _render_preference_item_tokens(rejected, renderer=renderer, tokenizer=tokenizer) + if len(chosen_tokens) < 2 or len(rejected_tokens) < 2: + return None + if max_seq_len is not None and ( + len(chosen_tokens) > max_seq_len or len(rejected_tokens) > max_seq_len + ): + return None + + return RenderedPreferencePair( + chosen_tokens=chosen_tokens, + rejected_tokens=rejected_tokens, + response_start=_common_prefix_length(chosen_tokens, rejected_tokens), + chosen_datum=build_next_token_datum(chosen_tokens), + rejected_datum=build_next_token_datum(rejected_tokens), + ) diff --git a/training/utils/validation.py b/training/utils/validation.py index 0d9c7e3b..8d1a01df 100644 --- a/training/utils/validation.py +++ b/training/utils/validation.py @@ -4,11 +4,45 @@ import logging -from fireworks.training.sdk.errors import DOCS_HOTLOAD, DOCS_API_KEYS, DOCS_DEPLOYMENTS, format_sdk_error +from fireworks.training.sdk import errors as sdk_errors from training.utils.config import InfraConfig, DeployConfig, ResumeConfig, HotloadConfig logger = logging.getLogger(__name__) +_DEFAULT_DOCS_HOTLOAD = "https://docs.fireworks.ai/fine-tuning/hotload" +_DEFAULT_DOCS_API_KEYS = "https://fireworks.ai/account/api-keys" +_DEFAULT_DOCS_DEPLOYMENTS = "https://docs.fireworks.ai/deployments" + + +def _get_sdk_docs_url(name: str, default: str) -> str: + value = getattr(sdk_errors, name, None) + return value if isinstance(value, str) and value else default + + +def _format_sdk_error( + what: str, + cause: str, + solution: str, + docs_url: str | None = None, +) -> str: + formatter = getattr(sdk_errors, "format_sdk_error", None) + if callable(formatter): + return formatter(what, cause, solution, docs_url=docs_url) + + lines = [ + f"ERROR: {what}", + f" Cause: {cause}", + f" Solution: {solution}", + ] + if docs_url: + lines.append(f" Docs: {docs_url}") + return "\n".join(lines) + + +DOCS_HOTLOAD = _get_sdk_docs_url("DOCS_HOTLOAD", _DEFAULT_DOCS_HOTLOAD) +DOCS_API_KEYS = _get_sdk_docs_url("DOCS_API_KEYS", _DEFAULT_DOCS_API_KEYS) +DOCS_DEPLOYMENTS = _get_sdk_docs_url("DOCS_DEPLOYMENTS", _DEFAULT_DOCS_DEPLOYMENTS) + def validate_config( base_model: str, @@ -23,7 +57,7 @@ def validate_config( if not base_model: errors.append( - format_sdk_error( + _format_sdk_error( "Missing base_model", "No base model specified.", "Set base_model (e.g. 'accounts/fireworks/models/qwen3-8b').", @@ -31,7 +65,7 @@ def validate_config( ) elif not base_model.startswith("accounts/"): errors.append( - format_sdk_error( + _format_sdk_error( "Invalid base_model format", f"'{base_model}' doesn't match expected format.", "Use format: accounts/ACCOUNT/models/MODEL_NAME\n" " Example: accounts/fireworks/models/qwen3-8b", @@ -40,7 +74,7 @@ def validate_config( if not dataset: errors.append( - format_sdk_error( + _format_sdk_error( "Missing dataset", "No dataset path or URL specified.", "Set dataset to a local path or URL to a JSONL file.", @@ -57,7 +91,7 @@ def validate_config( if infra and infra.node_count < 1: errors.append( - format_sdk_error( + _format_sdk_error( "Invalid node_count", f"node_count={infra.node_count} must be >= 1.", "Set InfraConfig(node_count=1) or higher.", @@ -81,7 +115,7 @@ def validate_preflight( if not skip_credential_check: if not fw_api_key: errors.append( - format_sdk_error( + _format_sdk_error( "Missing FIREWORKS_API_KEY", "No API key found in --fireworks-api-key or FIREWORKS_API_KEY env var.", "export FIREWORKS_API_KEY='your-key-here'\n" @@ -91,7 +125,7 @@ def validate_preflight( ) if not fw_account_id: errors.append( - format_sdk_error( + _format_sdk_error( "Missing FIREWORKS_ACCOUNT_ID", "No account ID found in --fireworks-account-id or FIREWORKS_ACCOUNT_ID env var.", "export FIREWORKS_ACCOUNT_ID='your-account-id'\n"