Transformers documentation
連続バッチ処理
連続バッチ処理
このページでは、連続バッチ推論の背後にあるクラスについて説明します。プロンプトの送信、スケジューリングとメモリ制限の設定、結果の取得です。
使用例については、Continuous batchingガイドを、スケジューリングとメモリがどのように相互作用するかについては、Continuous batching architectureドキュメントを参照してください。
ContinuousMixin.generate_batch
transformers.ContinuousMixin.generate_batch
< source >( inputs: listgeneration_config: transformers.generation.configuration_utils.GenerationConfig | None = Nonecontinuous_batching_config: transformers.generation.configuration_utils.ContinuousBatchingConfig | None = Nonerecord_timestamps: bool = Falseprogress_bar: bool = Falsepersistent_manager: bool = Falsewarmup: bool = True**kwargs ) → dict[str, GenerationOutput]
パラメータ
- inputs — 入力トークンシーケンス（プロンプト）のリスト
- generation_config — 任意の生成設定
- continuous_batching_config — 任意の連続バッチ処理設定
- record_timestamps — true に設定すると、生成された各トークンにタイムスタンプが記録されます
- progress_bar — true に設定すると、プログレスバーが表示されます
- persistent_manager — 生成が完了した後もマネージャーを保持するかどうか。デフォルトは False です。
- warmup — リクエストの処理前に CUDA グラフを事前キャプチャするかどうか。デフォルトは True です。
戻り値
dict[str, GenerationOutput]
リクエストIDからGenerationOutputオブジェクトへの辞書
連続バッチ処理を使用して、プロンプトのバッチに対してシーケンスを生成します。
ContinuousBatchingManager
class transformers.ContinuousBatchingManager
< source >( model: ProtoPretrainedModelgeneration_config: GenerationConfigcontinuous_batching_config: ContinuousBatchingConfigworkload_hints: transformers.generation.continuous_batching.utils.WorkloadHints | None = None )
生成リクエストの連続バッチ処理を管理するマネージャーです。生成リクエストの送信、結果の取得、バックグラウンド生成スレッドの管理を行うユーザーインターフェースを提供します。このクラスは直接作成せず、以下のいずれかのエントリポイント（ContinuousMixinミックスインのすべてのメソッド）を通じて作成してください。
init_continuous_batchingcontinuous_batching_context_managergenerate_batch
add_request
< source >( input_ids: listrequest_id: str | None = Nonemax_new_tokens: int | None = Nonestreaming: bool = Falserecord_timestamps: bool = Falseeos_token_id: int | list[int] | None = None**logit_processor_kwargs: Any ) → str | None
パラメータ
- input_ids — プロンプトとして使用する入力トークンID
- request_id — 任意のカスタムリクエストID（None の場合は自動生成されます）
- max_new_tokens — 生成する新しいトークンの最大数
- streaming — トークンが生成されるにつれてストリーミングで出力するかどうか
- record_timestamps — 生成された各トークンについてタイムスタンプを記録するかどうか
- eos_token_id — シーケンス終了トークンID
- logit_processor_kwargs — logits プロセッサ用のキーワード引数。
戻り値
str | None
プロセスがTPドライバーである場合はリクエストID、それ以外の場合はNone。
キューに新しい生成リクエストを追加します。プロセスが TP ドライバーでない場合、この操作は何もしません。
add_requests
< source >( inputs: listmax_new_tokens: int | None = Nonestreaming: bool = Falserecord_timestamps: bool = False**logit_processor_kwargs: Any )
add_request をバッチ処理してその ID を返すユーティリティ関数です。詳細はドキュメントを参照してください。
IDを指定してリクエストをキャンセルします。このメソッドがTPドライバー以外のプロセスから呼び出された場合、何もしません。マネージャーとやり取りするのはTPドライバーのプロセスのみです。
マネージャーを終了し、分散リソースを解放します。複数回呼び出しても安全です。このメソッドを呼び出した後、マネージャーを再起動することはできません。
出力キューから1つの結果を取得します。IDが指定されている場合、最初に一致するリクエストを返します。タイムアウトが指定されている場合、タイムアウト（秒）後にNoneを返します。
バックグラウンド生成スレッドが開始され、かつまだ生存している場合は True を返します。
バックグラウンドスレッドの完了を待機します。waitは、timeout引数（秒）で上限を設定できます。
…生成ループを一時停止するコンテキストマネージャーです。これにより、呼び出し元スレッドがモデルを使用でき、通常はインプレースでモデルを更新できます。スレッドはループが一時停止した後にのみこのコンテキストに入ります。コンテキストを抜けた時点でループが再開され、キャッシュと進行中のリクエストは保持されます。何も排空されず、リクエストが失われることもありません。複数のスレッドが同時に一時停止を保持でき、最後のスレッドが抜けた時点でループが再開されます。TPが有効な場合、すべてのランクがこのコンテキストに入らなければなりません。そうでない場合、他のランクが永久にハングします。
register_result_handler
< source >( request_id: strcallback: Callable )
結果の配信（ストリーミングまたは非ストリーミング）用のコールバックを登録します。
このコールバックは、このリクエストの結果が生成されるたびに、イベントループ上でcall_soon_threadsafeを介して呼び出されます。ストリーミングリクエストの場合はトークンごとに、非ストリーミングの場合は完了時のみ呼び出されます。リクエストが完了すると、ハンドラーは自動的にクリーンアップされます。
特定のrequest idに一致する結果を反復処理します（ブロッキング）。
再キューイングを伴う共有出力キューを使用します。高並行度でのサービングには、代わりにregister_result_handlerを使用してください。
バックグラウンド生成スレッドを開始します。
stop
< source >( block: bool = Truetimeout: float | None = Nonekeep_for_next_session: bool = Falsehard_stop: bool = False )
バックグラウンド生成スレッドを停止します。block フラグが True に設定されている場合、このメソッドはスレッドの停止を最大 timeout 秒間待機します（None はタイムアウトなしを意味します）。keep_for_next_session フラグが True に設定されている場合、マネージャーは将来の使用のためにモデル上でキャッシュされます。hard_stop フラグが設定されている場合、バックグラウンド生成スレッドは即座に停止され、保留中のリクエストは失敗します。
switch_to_cb_friendly_attn
< source >( model: ProtoPretrainedModelauto_switch_to_flash: bool = True )
attnの実装を、CB互換のものに切り替えます。auto_switch_to_flashがTrueの場合、かつattn実装がSDPAまたはeagerの場合、サポートおよび利用可能な場合はflashに切り替えます。
ダミーバッチを実行して、varlenパスおよびデコードパスのCUDAグラフを事前キャプチャします。バッチプロセッサがまだ初期化されていない場合は、初期化します。
連続バッチ処理の設定
class transformers.ContinuousBatchingConfig
< source >( page_size: int = 256num_blocks: int | None = Nonemax_batch_tokens: int | None = Nonemax_memory_percent: float | None = Nonemax_requests_per_batch: int | None = Nonemax_blocks_per_request: int | None = Noneallow_block_sharing: bool = Trueuse_async_batching: bool | None = Noneuse_cuda_graph: bool | tuple[bool, bool] | None = Noneq_padding_interval_size: int = 0kv_padding_interval_size: int = 0varlen_compile_config: transformers.generation.configuration_utils.CompileConfig | None = Nonedecode_compile_config: transformers.generation.configuration_utils.CompileConfig | None = Nonedefault_compile_level: int = 0scheduler_type: str = 'fifo'safety_margin: float | None = Nonereturn_logprobs: bool = Falseseed: int | None = Noneauto_switch_to_flash: bool = Truecpu_offload_space: float | None = 0.0cpu_offload_space_safety_threshold: float = 0.8max_queue_size: int = 0per_request_processors: bool = Falsedrop_unsupported_processors: bool = Truedisable_nccl_graph_mixing: bool = Truecpu_group_timeout: float | None = 300.0use_default_compile_configs: bool | None = Nonemax_cached_graphs: int | None = Noneblock_size: int | None = None )
パラメータ
- page_size (
int、任意、デフォルトは 256) — (full-attention) ページ内で各レイヤーに保存されるトークン数です。N 個のレイヤーのキャッシュを保存するブロックには N 個のページ（レイヤーごとに 1 個）があり、それぞれが 1 つのレイヤーのpage_sizeトークン分のキャッシュを保持します。デフォルトは 256 です。 - num_blocks (
int, 任意) — KVキャッシュのブロック数。Noneの場合、GPUメモリから自動推定されます。 - max_batch_tokens (
int, 省略可) — バッチ内の最大トークン数。Noneの場合、GPU メモリから自動推定されます。 - max_memory_percent (
float, 任意) — モデル読み込み後の空き GPU メモリにおいて、KV キャッシュに使用できる最大割合です。Noneの場合、実行時に、logit 処理がない場合は 0.9、ある場合は 0.8 に解決され、語彙サイズ分の一時テンソルに余剰容量を確保します。 - max_requests_per_batch (
int, 任意) — バッチあたりの最大リクエスト数。Noneの場合はワークロードのヒントから自動推論され、フォールバック値として 1024 が使用されます。 - max_blocks_per_request (
int, 任意) — 1リクエストあたりの最大ブロック数。flash_attn_with_kvcacheの高速デコードパスにおいて、ブロックテーブルの次元を決定するために使用されます。0に設定すると、高速デコードパスが無効になります。デフォルトは None（自動推定）です。 - allow_block_sharing (
bool、任意、デフォルトはTrue) — プレフィックスキャッシュのためにブロック共有を許可するかどうか。一部のモデルがブロック共有をサポートしていないため、ブロック共有は許可することはできても、強制することはできません。短いプロンプトが少なく、生成長が長い場合は無効にしてください。 - use_async_batching (
bool, 省略可) — 非同期ダブルバッファリングを有効にするかどうか。VRAM 使用量が 2 倍になる代わりに、連続バッチングループからの CPU オーバーヘッドを排除します。Noneの場合は自動検出されます。 - use_cuda_graph (
boolまたはtuple[bool, bool]、省略可) — CUDA グラフを有効にするかどうかを指定します。これは、ブール値のタプル（varlen パス用と decode fast path 用をそれぞれ1つずつ）、両方のパスに適用されるブール値、または None（自動推論）のいずれかです。decide_use_cuda_graphsを呼び出した後、この属性はブール値のタプルになります。デフォルトは None（自動推論）です。 - q_padding_interval_size (
int、任意、デフォルト値は 0)— CUDA グラフにおけるクエリパディングの粒度(トークン数)。0 に設定されている場合、continuous_api.pyのプリセットを使用します。 - kv_padding_interval_size (
int、任意、デフォルトは 0) — CUDA グラフにおける KV padding のトークン単位の粒度です。0 に設定されている場合、continuous_api.pyのプリセットを使用します。 - varlen_compile_config (
CompileConfig、任意) — varlen（prefill）パス用の CompileConfig です。デフォルトは None（generation_config のフォールバックを使用）です。 varlen パスでは、クエリ長と KV 長が異なるバッチを処理するため、dynamic=True を指定すると効果的な場合があります。 - decode_compile_config (
CompileConfig、省略可) — デコード（高速）パス用の CompileConfig。デフォルトは None（generation_config のフォールバックを使用） デコードパスではバッチが動的な KV 長を持たないため、静的な形状の方が適しています。 - default_compile_level (
int、任意、デフォルトは 0) — この値が 0 より大きく、varlen パスまたは decode パスにコンパイル設定が指定されていない場合、デフォルトのコンパイル設定が提供されます。レベルは最大 3 まで設定でき、レベルが高いほどパフォーマンスは向上しますが、ウォームアップ時間が長くなります。 - scheduler_type (
str、任意、デフォルトは"fifo") — 使用するスケジューラーの種類。 - safety_margin (
float、省略可) — オフロード量を制限するために使用される安全マージン。デフォルトは None(クラスのデフォルトを使用)。 - return_logprobs (
bool、省略可、デフォルトはFalse) — 生成されたトークンとともに対数確率を返すかどうかを指定します。 - seed (
int | None, 任意) — 生成用の任意のシードです。指定しない場合、内部のシードはランダムな値に設定されます。 - auto_switch_to_flash (
bool、任意、デフォルトはTrue) — 実装が SDPA または eager の場合、利用可能な場合に自動的に flash attention に切り替えるかどうか。デフォルトは True です。flash は continuous batching のために大幅に優れているためです。 - cpu_offload_space (
float、任意、デフォルト値は 0.0)— KV キャッシュのオフロードに使用する CPU スワップスペースのサイズ (GiB)。初期化時に、このサイズ分のピン留め済み CPU バッファが事前割り当てされます。GPU キャッシュが満杯になると、退避されたリクエストの KV キャッシュは破棄される代わりに、このバッファにコピーされます。0 を指定するとオフロードが無効になります (デフォルト)。 - cpu_offload_space_safety_threshold (
float、任意、デフォルト値は 0.8) —cpu_offload_spaceがシステム全体の RAM のこの割合を超えると、ホスト上の OOM を回避するために値がクランプされます。 1.0 に設定すると、この安全上限が無効になります。psutil が利用できない場合は無視されます。 - max_queue_size (
int、任意、デフォルトは 0)— サービング用の最大リクエストキューサイズです。0 は無制限を意味します。 - per_request_processors (
bool、任意、デフォルトはFalse) — リクエストごとの logits プロセッサのパラメータを有効にします。デフォルトは False です。 - drop_unsupported_processors (
bool、省略可、デフォルトはTrue) — サポートされていない logits プロセッサをエラーにせず削除します。デフォルトは True です。 - disable_nccl_graph_mixing (
bool、任意、デフォルトはTrue) — NCCL の並列グラフキャプチャされた通信に対するセーフティネットを無効にします。CB では発生せず、TP にパフォーマンスの向上をもたらします。 - cpu_group_timeout (
float、任意、デフォルト値は 300.0) — CPU 通信がタイムアウトし、プロセスがクラッシュするまでの時間（秒）です。タイムアウトを無効にする場合は None に設定してください。デフォルトは 300 秒です。 - use_default_compile_configs (
bool | None、任意) — 5.11 で非推奨: default_compile_level を使用してください。 - max_cached_graphs (
int, 省略可) — 5.13 で非推奨: グラフの最大数はもはや問題ではありません。 - block_size (
int | None、省略可) — 5.17 で非推奨: 代わりに page_size が使用されます。
generate_batch メソッドまたは continuous_batching_context_manager コンテキストマネージャーを通じて continuous batching を使用する場合、continuous batching に関する引数を持つクラス。