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/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.""" 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]: 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,