From 66afe8ac8a219cb2c89b3d0e028e3bcaa5e6e7e5 Mon Sep 17 00:00:00 2001 From: Qiyuan Chen Date: Mon, 2 Nov 2026 15:56:55 +0800 Subject: [PATCH 1/3] [v5] Fix resource estimates for composite models and the vLLM scheduler chunk Qwen3.5 keeps its shape on text_config, and its Q/K/V size is head_dim rather than hidden_size over the head count. max_num_batched_tokens is the prefill chunk, capped by max_model_len, so a profile run does not treat the whole KV cache as a vision batch. Co-authored-by: Cursor --- .../async_rl/cluster_adapter.py | 49 ++++++++++++++++--- 1 file changed, 42 insertions(+), 7 deletions(-) diff --git a/src/twinkle_agentic/async_rl/cluster_adapter.py b/src/twinkle_agentic/async_rl/cluster_adapter.py index f676ca8b0..a92b74a15 100644 --- a/src/twinkle_agentic/async_rl/cluster_adapter.py +++ b/src/twinkle_agentic/async_rl/cluster_adapter.py @@ -229,8 +229,15 @@ def model_memory(self, config: Any, *, model_parallel_width: int, quantization: hidden = int(getattr(config, 'hidden_size', 0) or getattr(config, 'd_model', 0) or 0) heads = int(getattr(config, 'num_attention_heads', 0) or 1) kv_heads = int(getattr(config, 'num_key_value_heads', 0) or heads) + raw_head_dim = getattr(config, 'head_dim', None) + try: + head_dim = int(raw_head_dim) + except (TypeError, ValueError): + head_dim = 0 + if head_dim <= 0: + head_dim = hidden // max(1, heads) kv_dtype_bytes = self._dtype_bytes(getattr(config, 'torch_dtype', None)) - kv_bytes = math.ceil(2 * layers * kv_heads * (hidden // heads) * kv_dtype_bytes / model_parallel_width) + kv_bytes = math.ceil(2 * layers * kv_heads * head_dim * kv_dtype_bytes / model_parallel_width) graph_reserve = 0 if enforce_eager else int(min(1.0, max(0.15, parameter_count / 1e11)) * _GIB) return ModelMemoryEstimate( parameter_count=parameter_count, @@ -329,12 +336,20 @@ def estimate(self, *, hf_config: Any, trainer_gpus: int, auxiliary_widths: Mappi kv_budget = budget - memory.weight_bytes - memory.runtime_reserve_bytes - memory.graph_reserve_bytes if kv_budget <= 0: raise ValueError('sampler KV budget is non-positive after weights and runtime reserves.') - max_batched = max(1, math.floor(0.95 * kv_budget / max(1, memory.kv_bytes_per_token))) - max_batched = int(overrides.get('max_num_batched_tokens', max_batched)) - if max_batched < 1: - raise ValueError('max_num_batched_tokens must be positive.') - max_seqs = max(1, max_batched // max(1, total_p95)) + # Tokens the KV cache can hold. The scheduler chunk is separate: one forward + # should cover a full max_model_len prefill, not the entire cache. Writing the + # cache size into max_num_batched_tokens also becomes vLLM's vision-encoder + # budget and profiles dozens of max-size images. + kv_tokens = max(1, math.floor(0.95 * kv_budget / max(1, memory.kv_bytes_per_token))) + context_limit = int(engine_context or 0) + if context_limit <= 0: + context_limit = int(getattr(hf_config, 'max_position_embeddings', 0) or 0) configured_seqs = getattr(rollout_config, 'vllm_max_num_seqs', None) if backend == 'vllm' else None + if 'max_num_batched_tokens' in overrides: + seq_budget = int(overrides['max_num_batched_tokens']) + else: + seq_budget = kv_tokens + max_seqs = max(1, seq_budget // max(1, total_p95)) if configured_seqs is not None: max_seqs = min(max_seqs, configured_seqs) max_seqs = int(overrides.get('max_num_seqs', max_seqs)) @@ -343,6 +358,18 @@ def estimate(self, *, hf_config: Any, trainer_gpus: int, auxiliary_widths: Mappi if configured_seqs is not None and max_seqs > configured_seqs: raise ValueError(f'max_num_seqs override {max_seqs} exceeds the configured engine limit ' f'{configured_seqs}.') + if 'max_num_batched_tokens' in overrides: + max_batched = int(overrides['max_num_batched_tokens']) + else: + # One prefill step may be as long as the context window. Also keep room for + # one decode token from every live sequence, which vLLM requires. + chunk = kv_tokens if context_limit <= 0 else min(kv_tokens, context_limit) + max_batched = min(kv_tokens, max(chunk, max_seqs)) + if max_batched < 1: + raise ValueError('max_num_batched_tokens must be positive.') + if max_batched < max_seqs: + raise ValueError(f'max_num_batched_tokens={max_batched} is smaller than max_num_seqs={max_seqs}; ' + 'one scheduler step must hold a decode token from every live sequence.') sampler_replicas = int(overrides.get('sampler_replicas', 1)) if sampler_replicas < 1: @@ -403,7 +430,15 @@ def _hf_config(model_config: Any) -> Any: kwargs['revision'] = model_config.model_revision if getattr(model_config, 'trust_remote_code', None) is not None: kwargs['trust_remote_code'] = model_config.trust_remote_code - return AutoConfig.from_pretrained(model_config.model, **kwargs) + config = AutoConfig.from_pretrained(model_config.model, **kwargs) + # Composite checkpoints (Qwen3.5) keep shape and dtype on text_config. The parent + # has neither hidden size nor vocab size, so estimation would refuse the model. + text = getattr(config, 'text_config', None) + if text is not None and int(getattr(text, 'hidden_size', 0) or 0) > 0: + if getattr(text, 'torch_dtype', None) is None and getattr(config, 'torch_dtype', None) is not None: + text.torch_dtype = config.torch_dtype + return text + return config def _prompt_lengths(template: Any, prompt_rows: Sequence[Mapping[str, Any]], limit: int = 2048) -> List[int]: From bcf7bfcc3357f9dc79e15aab0a6d4dab0963e90f Mon Sep 17 00:00:00 2001 From: Qiyuan Chen Date: Mon, 2 Nov 2026 15:56:55 +0800 Subject: [PATCH 2/3] [v5] Keep adaptive rollouts on the process that owns sampler admission configure_workers records live replicas on the driver handler. Pickling that handler drops the table and then refuses sample(), so MultiTurnRollout stays in this process once the pool is configured. Co-authored-by: Cursor --- src/twinkle/infra/dispatch.py | 5 +++++ src/twinkle_agentic/rollout/multi_turn.py | 13 +++++++++++++ 2 files changed, 18 insertions(+) diff --git a/src/twinkle/infra/dispatch.py b/src/twinkle/infra/dispatch.py index d0ee2d622..837a3d511 100644 --- a/src/twinkle/infra/dispatch.py +++ b/src/twinkle/infra/dispatch.py @@ -441,6 +441,11 @@ def _new_init_body(self, _caller, *args, **kwargs): caller_line = _cl # Pass an instance_id is recommended instance_id = kwargs.pop('instance_id', '') + f'{caller_file}_{caller_line}' + # A class may refuse to leave this process: its constructor arguments can own + # state that serialization drops (a sampler's admission table). + coerce = getattr(cls, '_coerce_remote_group', None) + if coerce is not None: + kwargs['remote_group'] = coerce(*args, **kwargs) remote_group = kwargs.get('remote_group') if os.environ.get('WORKER_NAME') is None and remote_group is None: logger.info(f'⚠️ Using local initialization of class: {cls}, please make sure the class ' diff --git a/src/twinkle_agentic/rollout/multi_turn.py b/src/twinkle_agentic/rollout/multi_turn.py index 1ea97a44f..6a45529fa 100644 --- a/src/twinkle_agentic/rollout/multi_turn.py +++ b/src/twinkle_agentic/rollout/multi_turn.py @@ -360,6 +360,19 @@ class MultiTurnRollout(Rollout): also enables sticky routing. Requires a continuous-work sampler. """ + @staticmethod + def _coerce_remote_group(*args, **kwargs): + """Keep the episode on the process that owns sampler admission. + + ``configure_workers`` records which replicas are ACTIVE on the driver handler. + Serialization drops that table and marks the copy so it cannot admit. ``sample()`` + therefore has to run here; a payload actor would be holding a handler that refuses. + """ + sampler = kwargs.get('sampler', args[0] if args else None) + if getattr(sampler, '_sampler_workers', None) is None: + return kwargs.get('remote_group') + return None + def __init__( self, sampler=None, From d20affd8a7ebab9b3c273c45853a5abbc0bbc58f Mon Sep 17 00:00:00 2001 From: Qiyuan Chen Date: Mon, 2 Nov 2026 15:56:56 +0800 Subject: [PATCH 3/3] [v5] Count live sample() calls when adaptive scaling checks load The current sample path holds dispatching for the whole call. The retired inflight ledger stays at zero, so the scaler never sees a busy sampler. Co-authored-by: Cursor --- src/twinkle/infra/worker_pool.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/twinkle/infra/worker_pool.py b/src/twinkle/infra/worker_pool.py index e953986f0..2bae0eb16 100644 --- a/src/twinkle/infra/worker_pool.py +++ b/src/twinkle/infra/worker_pool.py @@ -277,7 +277,15 @@ def worker_inflight(self, worker_ids): @property def inflight_count(self): - return sum(self.worker_inflight(self.worker_ids())) + """In-flight sample calls. + + ``inflight`` is the retired submit/collect ledger. The live ``sample()`` path + holds ``dispatching`` for the whole call, so a count of only ``inflight`` stays 0 + and the scaler never sees load. + """ + with self._worker_lock: + return sum(worker.inflight + worker.dispatching + for worker in self._sampler_workers.values()) def workers_idle(self, worker_ids): """Check control state only -- an actor is retired without reading generation output."""