maxent_grpo.training.rollout.vllm_colocate

In-process (colocated) vLLM generation helpers for the custom loop.

Functions

_apply_param_updates(index, updates[, ...])

_build_param_index(model)

_candidate_children(obj)

_coerce_logprob_payload(payload)

_configure_colocate_logging()

_dist_initialized()

_env_bool(name)

_env_float(name)

_env_int(name)

_extract_logprob_sequence(raw)

_filter_kwargs(callable_obj, kwargs)

_init_mode()

_log_env_snapshot(keys)

_log_process_snapshot()

_log_runtime_snapshot()

_log_torch_snapshot()

_lookup_param(name, index)

_outputs_to_payload(outputs, want_logprobs)

_parse_log_level(raw)

_reset_prefix_cache_llm(llm)

_resolve_dtype(ctx)

_resolve_llm_model(llm)

Return the first object in the vLLM stack exposing named_parameters().

_resolve_model_id(ctx)

_sum_logprobs(values)

_sync_chunk_bytes()

Return the max payload size for colocate sync batches (bytes).

_use_vllm_collective()

_vllm_colocate_worker(conn)

Subprocess worker for vLLM colocate init/generate.

Classes

ColocateVLLMClient(engine)

Local client adapter that mimics TRL's VLLMClient interface.

ColocateVLLMEngine(ctx, fallback_generate)

Lazy vLLM engine wrapper used for colocated generation.
class maxent_grpo.training.rollout.vllm_colocate.ColocateVLLMEngine(ctx, fallback_generate)[source]

Bases: object

Lazy vLLM engine wrapper used for colocated generation.

Parameters:

  • ctx (Any)

  • fallback_generate (Any)

sync_client()[source]
Return type:

ColocateVLLMClient

ensure_ready()[source]

Initialize the colocate worker/engine before parameter streaming begins.

Return type:

None

request_batch(prompts, request_count)[source]
Parameters:
Return type:

Tuple[List[List[str]] | None, List[List[VLLMLogprobResult | None]] | None]

class maxent_grpo.training.rollout.vllm_colocate.ColocateVLLMClient(engine)[source]

Bases: object

Local client adapter that mimics TRL’s VLLMClient interface.

Parameters:

engine (ColocateVLLMEngine)

update_named_param(name, param)[source]
Parameters:
Return type:

None

ensure_ready()[source]
Return type:

None

flush()[source]
Return type:

None

reset_prefix_cache()[source]
Return type:

None