feat: client-side embedding training & multi-turn rollout support

cookbook/client/twinkle/self_host/embedding.py

@@ -0,0 +1,121 @@
+# Twinkle Client - Embedding (InfoNCE contrastive) Training Example
+#
+# Client counterpart of the ray-local core-lib example
+# ``cookbook/exp/embedding/train_embedding_full_ddp.py`` and the client section
+# of ``docs/source_zh/使用指引/Embedding训练.md``. Instead of building a local
+# ``TransformersModel``/``MegatronModel`` + Ray, it drives the SAME training flow
+# over HTTP through the Twinkle client:
+#
+#   set_processor('InputProcessor')
+#     -> set_loss('InfonceLoss', ...)
+#     -> set_optimizer('Adam', ...)
+#     -> add_metric('EmbeddingMetric', is_training=True)
+#     -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step
+#     -> calculate_metric(is_training=True)
+#
+# Data format (anchor/positive contrastive pairs):
+#   inputs: [anchor_0, positive_0, anchor_1, positive_1, ...]
+#   labels:  [      1,          0,        1,          0, ...]
+#   (labels==1 marks a new group's anchor; labels==0 marks its positive)
+#
+# The server must be running first (see server.py / server_config.yaml). Only the
+# model service is required (no sampler). Works against both the transformers and
+# megatron backends: each single-sequence feature carries ``position_ids`` so the
+# megatron mrope model gets valid positions (transformers derives them internally).
+
+import dotenv
+
+dotenv.load_dotenv('.env')
+
+from typing import Any, Dict, List
+
+from peft import LoraConfig
+
+from twinkle import get_logger, init_twinkle_client
+from twinkle.template import Qwen3_5Template
+from twinkle_client.model import MultiLoraTransformersModel
+
+logger = get_logger()
+
+# ========== Configuration ==========
+MODEL_ID = 'ms://Qwen/Qwen3.5-4B'
+ADAPTER_NAME = 'emb_adapter'
+TEMPERATURE = 0.07
+LEARNING_RATE = 1e-4
+MAX_STEPS = 10
+MAX_LENGTH = 512
+
+# Small self-contained anchor/positive corpus. Each anchor's positive shares its
+# meaning, so InfoNCE has a clear contrastive signal to learn from.
+PAIRS: List[tuple] = [
+    ('What is the capital of France?', 'Paris is the capital city of France.'),
+    ('How do plants make their food?', 'Plants use photosynthesis to produce food from sunlight.'),
+    ('What is the boiling point of water?', 'Water boils at 100 degrees Celsius at sea level.'),
+    ('Who wrote the play Hamlet?', 'William Shakespeare wrote the tragedy Hamlet.'),
+    ('What is the speed of light?', 'Light travels at about 300,000 kilometers per second.'),
+    ('What language is spoken in Brazil?', 'The main language spoken in Brazil is Portuguese.'),
+]
+
+
+def build_minibatch(tokenizer) -> List[Dict[str, Any]]:
+    """Build one interleaved [anchor, positive, ...] embedding minibatch.
+
+    Each sample is a plain-dict feature with:
+      * ``input_ids``      -- tokenized text (plain Python list, JSON-serializable)
+      * ``attention_mask`` -- all ones
+      * ``labels``         -- a single scalar: 1 for an anchor, 0 for its positive
+      * ``position_ids``   -- ``range(len)`` so the megatron mrope model gets valid
+                              positions (transformers derives them internally)
+    """
+    minibatch: List[Dict[str, Any]] = []
+    for anchor_text, positive_text in PAIRS:
+        for text, is_anchor in ((anchor_text, True), (positive_text, False)):
+            input_ids = tokenizer.encode(text, add_special_tokens=True)
+            minibatch.append({
+                'input_ids': list(input_ids),
+                'attention_mask': [1] * len(input_ids),
+                'labels': [1 if is_anchor else 0],
+                'position_ids': list(range(len(input_ids))),
+            })
+    return minibatch
+
+
+def train():
+    # Step 1: connect to the running Twinkle server.
+    init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN')
+
+    # Step 2: build the client model with a fresh LoRA adapter.
+    model = MultiLoraTransformersModel(model_id=MODEL_ID)
+    model.add_adapter_to_model(ADAPTER_NAME, LoraConfig(target_modules='all-linear'))
+    model.set_template('Qwen3_5Template', model_id=MODEL_ID)
+
+    # Step 3: configure the embedding-training pipeline (ORDER MATTERS).
+    # Pass class names as strings; the server resolves them to core-lib classes.
+    model.set_processor('InputProcessor')
+    model.set_loss('InfonceLoss', temperature=TEMPERATURE, use_batch=True, hard_negatives=None)
+    # 'Adam' works on both transformers and megatron backends ('AdamW' is
+    # transformers-only; megatron expects 'Adam'/'default'/'MegatronOptimizer').
+    model.set_optimizer('Adam', lr=LEARNING_RATE)
+    model.add_metric('EmbeddingMetric', is_training=True)
+
+    # Step 4: build a small contrastive minibatch with a local tokenizer.
+    template = Qwen3_5Template(model_id=MODEL_ID, max_length=MAX_LENGTH, enable_thinking=False)
+    minibatch = build_minibatch(template.tokenizer)
+    logger.info(f'Built minibatch: {len(minibatch)} samples ({len(PAIRS)} anchor/positive pairs)')
+
+    # Step 5: training loop. task='embedding' selects the embedding pooling +
+    # InfoNCE path on the server (TransformersEmbeddingPatch/MegatronEmbeddingPatch
+    # is applied and rolled back automatically inside the forward).
+    for step in range(MAX_STEPS):
+        model.forward_backward(inputs=minibatch, task='embedding')
+        model.clip_grad_and_step(max_grad_norm=1.0)
+        metric = model.calculate_metric(is_training=True)
+        metric = getattr(metric, 'result', metric)
+        logger.info(f'Step {step}: {metric}')
+
+    logger.info('Embedding training finished. Healthy training shows pos_sim rising '
+                'and neg_sim stable/decreasing.')
+
+
+if __name__ == '__main__':
+    train()

cookbook/client/twinkle/self_host/multi_turn_rollout.py

@@ -0,0 +1,199 @@
+# Twinkle Client - Multi-turn agentic rollout Example
+#
+# Client counterpart of the ray-local core-lib example
+# ``cookbook/rl/multi_turn/multi_turn_grpo.py``. Instead of a Ray ``MultiTurnRollout``
+# + ``EnvPool``, it drives the SAME multi-turn "sample -> tool -> bridge -> sample"
+# loop over HTTP with ``twinkle_client.rollout.ClientMultiTurnRollout`` and the
+# client ``vLLMSampler``.
+#
+# Differences from the ray-local example (by design):
+#   * EnvPool is a Ray-only component and is NOT usable from the client, so this
+#     example uses a self-contained Python tool (a small calculator) wrapped in a
+#     ``ToolManager`` instead of an interactive environment.
+#   * ``ClientMultiTurnRollout`` samples with ``num_samples=1``; for a GRPO group we
+#     replicate each prompt ``NUM_GENERATIONS`` times as separate trajectories
+#     (exactly what the ray-local example does with ``BATCH_SIZE * NUM_GENERATIONS``).
+#
+# The server must be running first (see server.py / server_config.yaml). Both the
+# model and sampler services must be configured.
+#
+# NOTE on weight sync: ``ClientMultiTurnRollout`` samples with the sampler's
+# currently loaded server-side weights. A full RL loop would periodically
+# ``model.save(is_sampler=True)`` and point the sampler at the saved adapter; that
+# sync is intentionally omitted here to keep the rollout example focused.
+
+import dotenv
+
+dotenv.load_dotenv('.env')
+
+from typing import Any, Dict, List
+
+from peft import LoraConfig
+
+from twinkle import get_logger, init_twinkle_client
+from twinkle.advantage import GRPOAdvantage
+from twinkle.data_format import SamplingParams
+from twinkle.template import Qwen3_5Template
+from twinkle_agentic.tools.base import Tool
+from twinkle_agentic.tools.tool_manager import ToolManager
+from twinkle_client.model import MultiLoraTransformersModel
+from twinkle_client.rollout import ClientMultiTurnRollout
+from twinkle_client.sampler import vLLMSampler
+
+logger = get_logger()
+
+# ========== Configuration ==========
+MODEL_ID = 'ms://Qwen/Qwen3.5-4B'
+ADAPTER_NAME = 'default'
+NUM_GENERATIONS = 2       # GRPO group size (rollout runs num_samples=1 per trajectory)
+MAX_NEW_TOKENS = 512
+MAX_TURNS = 4
+LEARNING_RATE = 1e-5
+MAX_STEPS = 3
+
+SYSTEM_PROMPT = ('You are a careful assistant. When a calculation is needed, call the '
+                 '`calculator` tool with a single arithmetic ``expression`` (e.g. "12*7+3") '
+                 'and then give the final numeric answer.')
+
+# A few arithmetic questions; each is expanded into NUM_GENERATIONS trajectories.
+QUESTIONS: List[Dict[str, Any]] = [
+    {'q': 'What is 12 * 7 + 3?', 'answer': 87},
+    {'q': 'Compute (45 - 6) * 2.', 'answer': 78},
+]
+
+
+# ========== A self-contained tool (replaces EnvPool) ==========
+class CalculatorTool(Tool):
+    """Evaluate a single arithmetic ``expression`` string and return the result."""
+
+    _ALLOWED = set('0123456789+-*/(). ')
+
+    def tool_info(self) -> Dict[str, Any]:
+        return {
+            'type': 'function',
+            'function': {
+                'name': 'calculator',
+                'description': 'Evaluate one arithmetic expression and return its numeric value.',
+                'parameters': {
+                    'type': 'object',
+                    'properties': {
+                        'expression': {
+                            'type': 'string',
+                            'description': 'The arithmetic expression to evaluate, e.g. "12*7+3".',
+                        }
+                    },
+                    'required': ['expression'],
+                },
+            },
+        }
+
+    def __call__(self, tool_name: str, arguments: Dict[str, Any]) -> str:
+        expr = str(arguments.get('expression', '')).strip()
+        if not expr:
+            return 'Error: missing "expression".'
+        if not set(expr) <= self._ALLOWED:
+            return 'Error: expression may only contain digits, spaces and + - * / ( ) .'
+        try:
+            # Safe: the whitelist above forbids names, calls and attribute access.
+            return str(eval(expr, {'__builtins__': {}}, {}))  # noqa: S307
+        except Exception as e:  # noqa
+            return f'Error: could not evaluate {expr!r}: {type(e).__name__}'
+
+
+def build_trajectories(tool_schema: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """Expand each question into NUM_GENERATIONS identical trajectories (GRPO group)."""
+    trajectories: List[Dict[str, Any]] = []
+    for item in QUESTIONS:
+        for _ in range(NUM_GENERATIONS):
+            trajectories.append({
+                'messages': [
+                    {'role': 'system', 'content': SYSTEM_PROMPT},
+                    {'role': 'user', 'content': item['q']},
+                ],
+                'tools': tool_schema,
+            })
+    return trajectories
+
+
+def compute_rewards(trajectories: List[Dict[str, Any]]) -> List[float]:
+    """+1 if the correct numeric answer appears in the final assistant message."""
+    rewards: List[float] = []
+    n_per_q = NUM_GENERATIONS
+    for i, traj in enumerate(trajectories):
+        expected = str(QUESTIONS[i // n_per_q]['answer'])
+        final = ''
+        for msg in reversed(traj.get('messages', [])):
+            if msg.get('role') == 'assistant':
+                final = str(msg.get('content') or '')
+                break
+        rewards.append(1.0 if expected in final else 0.0)
+    return rewards
+
+
+def train():
+    # Step 1: connect to the running Twinkle server.
+    init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN')
+
+    # Step 2: training model (GRPO), mirroring the ray-local example's config.
+    model = MultiLoraTransformersModel(model_id=MODEL_ID)
+    model.add_adapter_to_model(ADAPTER_NAME, LoraConfig(target_modules='all-linear', r=16, lora_alpha=32))
+    model.set_loss('GRPOLoss', epsilon=0.2)
+    model.set_optimizer('Adam', lr=LEARNING_RATE)
+    model.set_processor('InputProcessor')
+    model.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False)
+
+    # Step 3: client sampler (HTTP).
+    sampler = vLLMSampler(model_id=MODEL_ID)
+    sampler.set_template('Qwen3_5Template', model_id=MODEL_ID, enable_thinking=False)
+
+    # Step 4: multi-turn rollout. Needs a LOCAL template for bridge-token stitching
+    # (rendering tool turns + the next generation prompt) and a ToolManager.
+    rollout_template = Qwen3_5Template(model_id=MODEL_ID, max_length=8192, enable_thinking=False)
+    rollout_template.truncation_strategy = 'delete'
+    tool_manager = ToolManager([CalculatorTool()])
+    tool_schema = tool_manager.tool_infos()
+
+    rollout = ClientMultiTurnRollout(
+        sampler=sampler,
+        template=rollout_template,
+        tool_manager=tool_manager,
+        sampling_params=SamplingParams(
+            max_tokens=MAX_NEW_TOKENS, num_samples=1, logprobs=1, temperature=1.0, top_p=0.95),
+        max_turns=MAX_TURNS,
+    )
+
+    advantage_fn = GRPOAdvantage()
+
+    for step in range(MAX_STEPS):
+        # 1. Run the batched multi-turn rollout (sample -> tool -> bridge -> sample).
+        trajectories = build_trajectories(tool_schema)
+        rolled = rollout(trajectories, tool_manager=tool_manager)
+
+        # 2. Read back per-trajectory logprobs (top-1 chosen-token logprob) and rewards.
+        all_inputs: List[Dict[str, Any]] = []
+        all_old_logps: List[List[float]] = []
+        for traj in rolled:
+            logprobs = traj.get('logprobs') or []
+            all_old_logps.append([lp[0][1] for lp in logprobs])
+            all_inputs.append(traj)
+        rewards = compute_rewards(rolled)
+
+        avg_turns = sum(t.get('turns') or 0 for t in rolled) / len(rolled)
+        logger.info(f'[Step {step}] avg_reward={sum(rewards) / len(rewards):.3f} avg_turns={avg_turns:.1f}')
+
+        # 3. GRPO advantages (group-relative within NUM_GENERATIONS).
+        advantages = advantage_fn(rewards, num_generations=NUM_GENERATIONS, scale='group').tolist()
+        if all(abs(a) < 1e-8 for a in advantages):
+            logger.info(f'[Step {step}] all advantages zero, skipping update')
+            continue
+
+        # 4. Policy update over the rolled-out trajectories.
+        model.forward_backward(inputs=all_inputs, advantages=advantages, old_logps=all_old_logps)
+        model.clip_grad_and_step()
+        logger.info(f'[Step {step}] {model.calculate_metric(is_training=True).result}')
+
+    logger.info('Multi-turn rollout example finished.')
+
+
+if __name__ == '__main__':
+    train()

docs/source_zh/使用指引/Embedding训练.md

@@ -106,6 +106,72 @@ for step, batch in enumerate(dataloader):
 
 ---
 
+## 通过 Twinkle Client 训练
+
+除了直接使用裸库 `TransformersModel`，你也可以通过 `twinkle_client` 的 `MultiLoraTransformersModel` 以 HTTP 方式驱动服务端完成 Embedding 训练。协议层（`/twinkle/*`）已经具备完成 Embedding 训练所需的全部能力，**无需引入任何新的高层封装函数**（例如 `setup_embedding_training`），只需按正确顺序调用现有方法即可。
+
+### 完整调用顺序
+
+client 化的 Embedding 训练遵循与裸库一致的调用顺序：
+
+```
+set_processor('InputProcessor')
+  -> set_loss('InfonceLoss', ...)
+  -> add_metric('EmbeddingMetric', is_training=True)
+  -> loop: forward_backward(inputs=mb, task='embedding') + clip_grad_and_step(...)
+  -> calculate_metric(is_training=True)
+```
+
+与裸库不同的是，client 侧传入的是**类名字符串**（如 `'InfonceLoss'`、`'EmbeddingMetric'`、`'InputProcessor'`），由服务端解析为对应的核心库类。
+
+### 示例
+
+```python
+from peft import LoraConfig
+from twinkle_client import init_twinkle_client
+from twinkle_client.model import MultiLoraTransformersModel
+
+# --- Connect to the running Twinkle server ---
+init_twinkle_client(base_url='http://127.0.0.1:8000', api_key='EMPTY_TOKEN')
+
+# --- Build the client model with a LoRA adapter for embedding ---
+model = MultiLoraTransformersModel(model_id='ms://Qwen/Qwen3.5-4B')
+model.add_adapter_to_model('emb_adapter', LoraConfig(target_modules='all-linear'))
+model.set_template('Qwen3_5Template')
+
+# --- Configure the embedding-training pipeline (order matters) ---
+# NOTE: pass class names as strings; the server resolves them to core-lib classes.
+model.set_processor('InputProcessor')
+model.set_loss('InfonceLoss', temperature=0.07, use_batch=True, hard_negatives=None)
+model.add_metric('EmbeddingMetric', is_training=True)
+
+# --- Training loop ---
+for step, mb in enumerate(minibatches):
+    # `task='embedding'` is forwarded through the /twinkle/* protocol as an extra
+    # kwarg and selects the embedding pooling + InfoNCE loss path on the server.
+    model.forward_backward(inputs=mb, task='embedding')
+    model.clip_grad_and_step(max_grad_norm=1.0)
+
+# --- Read back embedding metrics (pos_sim / neg_sim / loss) ---
+metric = model.calculate_metric(is_training=True)
+```
+
+### 关于 `TransformersEmbeddingPatch` 的自动应用
+
+Embedding 训练依赖 `TransformersEmbeddingPatch` 把 causal LM 的 `lm_head` 替换为 identity 输出，从而得到 per-token hidden states 用于池化。**你不需要显式调用 `apply_patch(TransformersEmbeddingPatch())`**：当你在 `forward_backward`（或 `forward_only`）中传入 `task='embedding'` 时，服务端的 `_resolve_task_context` 会在该次 `forward` 调用期间自动应用该 patch，并在调用结束后自动回滚。
+
+这意味着：
+
+- 同一个 adapter 在一次 `forward_backward(task='embedding')` 之后，紧接着调用不带 `task` 参数的 `forward_only` 仍会返回正常的词表维度 `logits`，不会残留上一次 embedding 任务的 identity hidden states。
+- 你在 client 侧真正需要显式调用的前置步骤只有三个：`set_processor('InputProcessor')`、`set_loss('InfonceLoss', ...)` 与 `add_metric('EmbeddingMetric', is_training=True)`。
+
+### 说明
+
+- 本节不引入任何新的高层封装函数，仅说明现有 `MultiLoraTransformersModel` 方法的正确调用顺序。
+- 各参数（`temperature`、`use_batch`、`hard_negatives` 等）的含义与取值建议参见上文「关键参数」；`calculate_metric` 返回的指标含义参见下文「监控指标」。
+
+---
+
 ## 监控指标
 
 `EmbeddingMetric` 报告关键训练信号：