Accelerate documentation

大規模モデルの扱い

Accelerate

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v1.15.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

大規模モデルの扱い

ディスパッチとオフロード

init_empty_weights

accelerate.init_empty_weights

< source >

( include_buffers: bool | None = None )

パラメータ

include_buffers (bool, 任意) — 初期化時に、すべてのバッファもメタデバイスに配置するかどうかを指定します。

モデルのパラメータをすべてメタデバイス上で初期化し、空のモデルを作成するコンテキストマネージャーです。モデルの初期化だけで利用可能なRAMが不足する場合に便利です。

例:

import torch.nn as nn
from accelerate import init_empty_weights

# Initialize a model with 100 billions parameters in no time and without using any RAM.
with init_empty_weights():
    tst = nn.Sequential(*[nn.Linear(10000, 10000) for _ in range(1000)])

Any model created under this context manager has no weights. As such you can’t do something like model.to(some_device) with it. To load weights inside your empty model, see load_checkpoint_and_dispatch(). Make sure to overwrite the default device_map param for load_checkpoint_and_dispatch(), otherwise dispatch is not called.

cpu_offload

accelerate.cpu_offload

< source >

( model: Moduleexecution_device: torch.device | None = Noneoffload_buffers: bool = Falsestate_dict: dict[str, torch.Tensor] | None = Nonepreload_module_classes: list[str] | None = None )

パラメータ

model (torch.nn.Module) — オフロードするモデル。
execution_device (torch.device, 任意) — モデルのフォワードパスが実行されるデバイス（GPUである必要があります）。デフォルトでは、モデルの最初のパラメータのデバイスが使用されます。
offload_buffers (bool、任意、デフォルトはFalse) — モデルのパラメータとともにバッファをオフロードするかどうか。
state_dict (Dict[str, torch.Tensor], 省略可) — CPU に保持されるモデルの state dict です。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。

モデルのフルCPUオフロードを有効にします。これにより、モデルのすべてのパラメータがオフロードされ、モデルのstate dictのコピーは1つだけ保持されます。フォワードパスの間、パラメータはそのstate dictから抽出され、必要に応じて実行デバイスにそのまま配置され、その後再度オフロードされます。

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

< source >

( model: Moduleexecution_device: int | str | torch.device | None = Noneprev_module_hook: accelerate.hooks.UserCpuOffloadHook | None = None )

パラメータ

model (torch.nn.Module) — オフロードするモデル。
execution_device(str, int または torch.device, optional) — モデルを実行するデバイスです。利用可能な場合は MPS デバイスがデフォルトとなり、次にアクセラレータデバイスがある場合はデバイス 0、最後に CPU がデフォルトとなります。
prev_module_hook (UserCpuOffloadHook、任意) — 実行中のパイプライン内の、このフックがアタッチされるモデルより前のモデルに対して、この関数から返されるフックです。指定された場合、このフックがアタッチされるモデルの forward が実行される直前に、そのオフロードメソッドが呼び出されます。

モデルをCPUにオフロードし、実行時に実行デバイスに戻します。cpu_offload()との違いは、フォワード実行後もモデルが実行デバイス上に保持され、返されるhookのoffloadメソッドが呼び出されたときにのみ再度オフロードされる点です。ループ内でモデルを実行するパイプラインに有用です。

例:

model_1, hook_1 = cpu_offload_with_hook(model_1, device)
model_2, hook_2 = cpu_offload_with_hook(model_2, device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, device, prev_module_hook=hook_2)

hid_1 = model_1(input)
for i in range(50):
    # model1 is offloaded on the CPU at the first iteration, model 2 stays on the GPU for this whole loop.
    hid_2 = model_2(hid_1)
# model2 is offloaded to the CPU just before this forward.
hid_3 = model_3(hid_3)

# For model3, you need to manually call the hook offload method.
hook_3.offload()

disk_offload

accelerate.disk_offload

< source >

( model: Moduleoffload_dir: str | os.PathLikeexecution_device: torch.device | None = Noneoffload_buffers: bool = Falsepreload_module_classes: list[str] | None = None )

パラメータ

model (torch.nn.Module) — オフロードするモデル。
offload_dir (str または os.PathLike) — モデルの重みをオフロードするフォルダ（または、すでにオフロードされているフォルダ）。
execution_device (torch.device, 任意) — モデルのフォワードパスが実行されるデバイス（GPUである必要があります）。デフォルトでは、モデルの最初のパラメータのデバイスが使用されます。
offload_buffers (bool、任意、デフォルトはFalse) — モデルのパラメータとともにバッファをオフロードするかどうか。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。

モデルのフルディスクオフロードを有効にします。これにより、モデルのすべてのパラメータが、指定されたフォルダー内のメモリマップド配列としてオフロードされます。フォワードパスの間、パラメータは必要に応じてそのフォルダーからアクセスされ、実行デバイスにそのまま配置され、その後再びオフロードされます。

dispatch_model

accelerate.dispatch_model

< source >

パラメータ

model (torch.nn.Module) — ディスパッチするモデル。
device_map (Dict[str, Union[str, int, torch.device]]) — モデルのstate_dict内のモジュール名を、配置先デバイスにマッピングする辞書です。"disk"はtorch.deviceの正式な値ではありませんが、受け入れられます。
main_device (str、int または torch.device、省略可) — メインの実行デバイスです。device_map 内の "cpu" または "disk" 以外の最初のデバイスがデフォルトになります。
state_dict (Dict[str, torch.Tensor], 任意) — CPU に保持されるモデルの一部分の state dict です。
offload_dir (str または os.PathLike) — モデルの重みをオフロードするフォルダ（または、すでにオフロードされているフォルダ）。
offload_index (Dict, 省略可) — 重み名からその情報(dtype/ shapeまたはsafetensorsファイル名)への辞書です。デフォルトでは、save_folderに保存されたインデックスが使用されます。
offload_buffers (bool、任意、デフォルトはFalse) — モデルのパラメータとともにバッファをオフロードするかどうか。
skip_keys (str または List[str]、省略可能) — デバイス間で入力または出力を移動する際に無視するキーのリストです。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。
force_hooks (bool、省略可、デフォルトはFalse) — すべてのレイヤーが単一のデバイスにディスパッチされている場合でも、モデルにデバイスフックを強制適用するかどうか。

与えられたデバイスマップに従ってモデルをディスパッチします。モデルのレイヤーは、複数のGPUに分散配置される場合もあれば、CPUやディスクにオフロードされる場合もあります。

load_checkpoint_and_dispatch

accelerate.load_checkpoint_and_dispatch

< source >

パラメータ

model (torch.nn.Module) — チェックポイントを読み込みたいモデル。
checkpoint (str または os.PathLike) — 読み込むフォルダのチェックポイントです。以下の場合があります:
- モデル全体の state dict を含むファイルへのパス
- シャード化されたチェックポイントのインデックスを含む.jsonファイルへのパス
- 一意の.index.jsonファイルとチェックポイントのシャードを含むフォルダへのパス。
device_map (Dict[str, Union[int, str, torch.device]], 省略可) — 各サブモジュールをどのデバイスに配置するかを指定するマップです。各パラメータ/バッファ名まで細かく指定する必要はなく、特定のモジュール名が指定されると、そのモジュールのすべてのサブモジュールが同じデバイスに送信されます。
Accelerate が最も最適化された device_map を自動的に計算するようにするには、device_map="auto" を設定します。各オプションの詳細については、こちらを参照してください。デフォルトは None で、dispatch_model() が呼び出されないことを意味します。
max_memory (Dict, 任意) — デバイス識別子から最大メモリ量への辞書です。未設定の場合、各 GPU の利用可能な最大メモリと、利用可能な CPU RAM がデフォルトとして使用されます。
no_split_module_classes (List[str], 任意) — デバイス間で分割すべきでないレイヤーのクラス名のリストです（残差接続を持つレイヤーなど）。
offload_folder (str または os.PathLike、省略可) — device_map に "disk" という値が含まれている場合、重みをオフロードするフォルダ。
offload_buffers (bool、任意、デフォルトはFalse) — CPUまたはハードドライブにオフロードされるレイヤーにおいて、パラメータに加えてバッファもオフロードするかどうか。
dtype (str または torch.dtype、省略可) — 指定された場合、読み込み時に重みがその型に変換されます。
offload_state_dict (bool, 省略可) — Trueの場合、CPU RAMを使い切らないように、CPU state dictの重みと最大のシャードの合計が収まらない場合に、CPU state dictを一時的にハードドライブへオフロードします。デバイスマップに"disk"の値が含まれている場合、デフォルトでTrueになります。
skip_keys (str または List[str]、省略可能) — デバイス間で入力または出力を移動する際に無視するキーのリストです。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。
force_hooks (bool、省略可、デフォルトはFalse) — すべてのレイヤーが単一のデバイスにディスパッチされている場合でも、モデルにデバイスフックを強制適用するかどうか。
strict (bool、省略可、デフォルトはFalse) — チェックポイントのstate_dictのキーが、モデルのstate_dictのキーと厳密に一致することを強制するかどうか。
full_state_dict (bool、省略可能、デフォルトはTrue) — Trueに設定すると、読み込まれた state_dict 内のすべてのテンソルが集約されます。読み込まれた state_dict には、ShardedTensor や DTensor は含まれません。
broadcast_from_rank0 (False、任意、デフォルトは False) — このオプションが True の場合、分散 ProcessGroup が初期化されている必要があります。rank0 は完全な state_dict を受信し、state_dict 内のテンソルを 1 つずつ他の rank にブロードキャストします。他の rank はテンソルを受信し、モデル内のローカルシャードに応じてシャード化（該当する場合）します。

モデル内の（シャード化された可能性のある）チェックポイントをロードし、ロードされる際に重みを指定されたデバイスに送信する場合があります。また、このモデルが正しく実行されるようにするためのさまざまなフックを追加します（デバイス間で分割されている場合でも）。

例:

>>> from accelerate import init_empty_weights, load_checkpoint_and_dispatch
>>> from huggingface_hub import hf_hub_download
>>> from transformers import AutoConfig, AutoModelForCausalLM

>>> # Download the Weights
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> weights_location = hf_hub_download(checkpoint, "pytorch_model.bin")

>>> # Create a model and initialize it with empty weights
>>> config = AutoConfig.from_pretrained(checkpoint)
>>> with init_empty_weights():
...     model = AutoModelForCausalLM.from_config(config)

>>> # Load the checkpoint and dispatch it to the right devices
>>> model = load_checkpoint_and_dispatch(
...     model, weights_location, device_map="auto", no_split_module_classes=["GPTJBlock"]
... )

load_checkpoint_in_model

accelerate.load_checkpoint_in_model

< source >

パラメータ

model (torch.nn.Module) — チェックポイントを読み込みたいモデル。
checkpoint (str または os.PathLike) — 読み込むフォルダのチェックポイントです。以下の場合があります:
- モデル全体の state dict を含むファイルへのパス
- シャード化されたチェックポイントのインデックスを含む.jsonファイルへのパス
- 一意の.index.jsonファイルとチェックポイントのシャードを含むフォルダへのパス。
- 一意のpytorch_model.binまたはmodel.safetensorsファイルを含むフォルダへのパス。
device_map (Dict[str, Union[int, str, torch.device]], 省略可) — 各サブモジュールをどのデバイスに配置するかを指定するマップです。パラメータやバッファの名前まで細かく指定する必要はなく、特定のモジュール名が指定されると、そのモジュールのすべてのサブモジュールが同じデバイスに送信されます。
offload_folder (str または os.PathLike、省略可) — device_map に "disk" という値が含まれている場合、重みをオフロードするフォルダ。
dtype (str または torch.dtype、省略可) — 指定された場合、読み込み時に重みがその型に変換されます。
offload_state_dict (bool、省略可、デフォルトはFalse) — Trueの場合、CPU state dictの重みと最大のシャードの合計がCPU RAMに収まらない場合に、CPU RAM不足を回避するため、CPU state dictを一時的にハードドライブにオフロードします。
offload_buffers (bool、任意、デフォルトはFalse) — ディスクにオフロードされる重みにバッファを含めるかどうかを指定します。
keep_in_fp32_modules(List[str], optional) — torch.float32 dtype で保持するモジュールのリスト。
offload_8bit_bnb (bool, 省略可) — 8ビットモジュールのCPU/ディスクへのオフロードを有効にするかどうか。
strict (bool、省略可、デフォルトはFalse) — チェックポイントのstate_dictのキーが、モデルのstate_dictのキーと厳密に一致することを強制するかどうか。
full_state_dict (bool、省略可能、デフォルトはTrue) — Trueに設定すると、読み込まれた state_dict 内のすべてのテンソルが集約されます。読み込まれた state_dict には、ShardedTensor や DTensor は含まれません。
broadcast_from_rank0 (False、任意、デフォルトは False) — このオプションが True の場合、分散 ProcessGroup が初期化されている必要があります。rank0 は完全な state_dict を受信し、state_dict 内のテンソルを 1 つずつ他の rank にブロードキャストします。他の rank はテンソルを受信し、モデル内のローカルシャードに応じてシャード化（該当する場合）します。

モデル内の（シャード化された可能性のある）チェックポイントをロードし、ロードされる際に重みを指定されたデバイスに送信する可能性があります。

Once loaded across devices, you still need to call dispatch_model() on your model to make it able to run. To group the checkpoint loading and dispatch in one single call, use load_checkpoint_and_dispatch().

infer_auto_device_map

accelerate.infer_auto_device_map

< source >

パラメータ

model (torch.nn.Module) — 解析対象のモデル。
max_memory (Dict, 任意) — デバイス識別子から最大メモリ量への辞書です。未設定の場合は、利用可能な最大メモリがデフォルトとして使用されます。例: max_memory={0: "1GB"}。
no_split_module_classes (List[str], 任意) — デバイス間で分割すべきでないレイヤーのクラス名のリストです（残差接続を持つレイヤーなど）。
dtype (str または torch.dtype、省略可) — 指定された場合、読み込み時に重みがその型に変換されます。
special_dtypes (Dict[str, Union[str, torch.device]], 省略可) — 指定した場合、特定の重みに適用する特殊な dtype です（すべての重みにデフォルトとして使用される dtype を上書きします）。
verbose (bool, optional, defaults to False) — Whether or not to provide debugging statements as the function builds the device_map.
clean_result (bool、省略可能、デフォルトはTrue) — 同じデバイスに配置されるすべてのサブモジュールをグループ化して、結果の device_map を整理します。
offload_buffers (bool、任意、デフォルトはFalse) — CPUまたはハードドライブにオフロードされるレイヤーにおいて、パラメータに加えてバッファもオフロードするかどうか。
fallback_allocation (bool、任意、デフォルトはFalse) — 通常の割り当てが失敗した場合、BFS を使用してサイズ制限内に収まるモジュールの割り当てを試みます。

…

どの GPU の利用可能なメモリも超過しません。
CPUへのオフロードが必要な場合、GPU 0には常に空き領域が残っており、CPUにオフロードされた層のうち最もサイズが大きいものを再度配置できます。
CPUへのオフロードが必要な場合、CPUで利用可能なRAMを超えないようにします。
ディスクへのオフロードが必要な場合、CPUには常に空き領域が残っており、ディスクにオフロードされたレイヤーのうち最もサイズが大きいものを再度配置できます。

すべての計算は、モデルパラメータのサイズとdtypeを解析することによって行われます。その結果、モデルはmetaデバイス上に配置できます（init_empty_weightsコンテキストマネージャー内で初期化された場合と同様です）。

フック

ModelHook

クラス accelerate.hooks.ModelHook

< source >

( )

モデルのforwardメソッドの直前と直後に実行されるコールバックを含むフックです。PyTorchの既存のフックとの違いは、kwargsが渡される点です。

クラス属性:

no_grad (bool, オプション, デフォルトは False) — torch.no_grad() コンテキストマネージャの下で実際のフォワードパスを実行するかどうか。

detach_hook

< source >

( module )

パラメータ

module (torch.nn.Module) — このフックから切り離されたモジュール。

フックがモジュールから切り離されたときに実行されます。

init_hook

< source >

( module )

パラメータ

module (torch.nn.Module) — このフックにアタッチされたモジュール。

フックがモジュールにアタッチされた際に実行されます。

post_forward

< source >

( moduleoutput ) → Any

パラメータ

module (torch.nn.Module) — このイベントの直前にフォワードパスが実行されたモジュール。
output (Any) — モジュールの出力。

戻り値

Any

処理されたoutput。

モデルの forward メソッドの直後に実行されます。

pre_forward

< source >

( module*args**kwargs ) → Tuple[Tuple[Any], Dict[Str, Any]]

パラメータ

module (torch.nn.Module) — このイベントの直後にフォワードパスが実行されるモジュール。
args (Tuple[Any]) — モジュールに渡される位置引数です。
kwargs (Dict[Str, Any]) — モジュールに渡されるキーワード引数。

戻り値

Tuple[Tuple[Any], Dict[Str, Any]]

処理されたargsとkwargsのタプル。

モデルの forward メソッドが実行される直前に実行されます。

AlignDevicesHook

class accelerate.hooks.AlignDevicesHook

< source >

パラメータ

execution_device (torch.device, 任意) — フォワードパスの前に、入力とモデルの重みを配置するデバイス。
offload (bool、省略可、デフォルトはFalse) — フォワードパスの後に重みをオフロードするかどうか。
io_same_device (bool、省略可能、デフォルトはFalse) — 出力を入力と同じデバイスに配置するかどうかを指定します。
weights_map (Mapping[str, torch.Tensor]、省略可) — モデルの重みがオフロードされている場合、パラメータ名からテンソル値への（遅延評価可能な）マップです。
offload_buffers (bool、任意、デフォルトはFalse) — オフロード時に、関連するモジュールのバッファを含めるかどうかを指定します。
place_submodules (bool、省略可能、デフォルトはFalse) — init_hookイベントの間に、サブモジュールをexecution_deviceに配置するかどうかを指定します。

関連するモジュールのフォワードパスにおいて、入力とモデルの重みが同じデバイスにあることを保証する汎用的なModelHook。フォワードパスの後に重みをオフロードする可能性があります。

SequentialHook

class accelerate.hooks.SequentialHook

< source >

( *hooks )

複数のフックを含み、各イベント時にそれらを順に処理するフック。

LayerwiseCastingHook

クラス accelerate.hooks.LayerwiseCastingHook

< source >

( storage_dtype: dtypecompute_dtype: dtypenon_blocking: bool )

モジュールの重みを計算用に高精度のdtypeにキャストし、保存用に低精度のdtypeにキャストするフックです。この処理により出力の品質が低下する可能性がありますが、メモリフットプリントを大幅に削減できます。

フックの追加

add_hook_to_module

accelerate.hooks.add_hook_to_module

< source >

( module: Modulehook: ModelHookappend: bool = False ) → torch.nn.Module

パラメータ

module (torch.nn.Module) — フックをアタッチするモジュール。
hook (ModelHook) — アタッチするフック。
append (bool、省略可能、デフォルトはFalse) — モジュールに既にフックが含まれている場合、既存のフックとチェーンするか、しないかを指定します。

戻り値

torch.nn.Module

同じモジュールにフックを付加した状態です（モジュールはインプレースで変更されるため、結果は破棄できます）。

指定されたモジュールにフックを追加します。これにより、モジュールのforwardメソッドがフックを含むように書き換えられます。この動作を削除して元のforwardメソッドを復元するには、remove_hook_from_moduleを使用します。

モジュールに既にフックが含まれている場合、デフォルトでは新しいフックで置き換えられます。2つのフックをチェーンするには、append=Trueを渡します。これにより、現在のフックと新しいフックがSequentialHookクラスのインスタンスとしてチェーンされます。

attach_execution_device_hook

accelerate.hooks.attach_execution_device_hook

< source >

パラメータ

module (torch.nn.Module) — フックを付与したいモジュール。
execution_device (int、str または torch.device) — フォワードパスの前に、入力とモデルの重みを配置するデバイス。
skip_keys (str または List[str]、省略可能) — デバイス間で入力または出力を移動する際に無視するキーのリストです。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 任意、デフォルトは None) — データポインタから、デバイスへの辞書（すでにディスパッチされた共有重み）へのマップです。特定の実行デバイスにおいて、このパラメータは、メモリを重複させる代わりに、共有重みの最初の利用可能なポインタを他のすべての重みで再利用するために役立ちます。

AlignDevicesHookを指定されたモデルのすべてのサブモジュールに再帰的にアタッチし、それらが正しい実行デバイスを持っていることを保証します

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

< source >

( module: Moduleexecution_device: torch.device | None = Noneoffload: bool = Falseweights_map: collections.abc.Mapping | None = Noneoffload_buffers: bool = Falsemodule_name: str = ''skip_keys: str | list[str] | None = Nonepreload_module_classes: list[str] | None = Nonetied_params_map: dict[int, dict[torch.device, torch.Tensor]] | None = None )

パラメータ

module (torch.nn.Module) — フックを付与したいモジュール。
execution_device (torch.device, 任意) — フォワードパスの前に、入力とモデルの重みを配置するデバイス。
offload (bool、省略可、デフォルトはFalse) — フォワードパスの後に重みをオフロードするかどうか。
weights_map (Mapping[str, torch.Tensor]、省略可) — モデルの重みがオフロードされている場合、パラメータ名からテンソル値への（遅延評価可能な）マップです。
offload_buffers (bool、任意、デフォルトはFalse) — オフロード時に、関連するモジュールのバッファを含めるかどうかを指定します。
module_name (str、任意、デフォルトは"") — モジュールの名前。
skip_keys (str または List[str]、省略可能) — デバイス間で入力または出力を移動する際に無視するキーのリストです。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 任意、デフォルトは None) — データポインタから、デバイスへの辞書（すでにディスパッチされた共有重み）へのマップです。特定の実行デバイスにおいて、このパラメータは、メモリを重複させる代わりに、共有重みの最初の利用可能なポインタを他のすべての重みで再利用するために役立ちます。

AlignDevicesHookを、直接のパラメータおよび/またはバッファを持つ指定モデルのすべてのサブモジュールに再帰的にアタッチします。

attach_align_device_hook_on_blocks

accelerate.hooks.attach_align_device_hook_on_blocks

< source >

パラメータ

module (torch.nn.Module) — フックを付与したいモジュール。
execution_device (torch.device または Dict[str, torch.device]、省略可) — フォワードパスの前に、入力とモデルの重みを配置するデバイスです。モジュール全体で1つのデバイスを使用することも、モジュール名からデバイスへのマッピングを辞書として指定することもできます。
offload (bool、省略可、デフォルトはFalse) — フォワードパスの後に重みをオフロードするかどうかを指定します。モジュール全体に対して1つのブール値を指定するか、モジュール名をブール値にマッピングする辞書を指定できます。
weights_map (Mapping[str, torch.Tensor]、省略可) — モデルの重みがオフロードされている場合、パラメータ名からテンソル値への（遅延評価可能な）マップです。
offload_buffers (bool、任意、デフォルトはFalse) — オフロード時に、関連するモジュールのバッファを含めるかどうかを指定します。
module_name (str、任意、デフォルトは"") — モジュールの名前。
skip_keys (str または List[str]、省略可能) — デバイス間で入力または出力を移動する際に無視するキーのリストです。
preload_module_classes (List[str]、省略可) — フォワードの開始時に、すべての重み（サブモジュール内のものを含む）を読み込むべきクラスのインスタンスのリストです。これは、フォワード中に直接呼び出されないが登録されているサブモジュールを持つクラスにのみ使用すべきです。例えば、dense 線形層が登録されているが、フォワードでは dense を直接呼び出す代わりに、dense.weight と dense.bias が一部の操作で使用される場合などが該当します。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 任意、デフォルトは None) — データポインタから、デバイスへの辞書（すでにディスパッチされた共有重み）へのマップです。特定の実行デバイスにおいて、このパラメータは、メモリを重複させる代わりに、共有重みの最初の利用可能なポインタを他のすべての重みで再利用するために役立ちます。

必要に応じて、指定されたモデルのすべてのブロックにAlignDevicesHookをアタッチします。

attach_layerwise_casting_hooks

accelerate.big_modeling.attach_layerwise_casting_hooks

< source >

( module: Modulestorage_dtype: dtypecompute_dtype: dtypeskip_modules_pattern: str | tuple[str, ...] | None = Noneskip_modules_classes: tuple[type[torch.nn.Module], ...] | None = Nonenon_blocking: bool = False )

パラメータ

module (torch.nn.Module) — 計算には高精度の dtype、保存には低精度の dtype にキャストされるリーフモジュールを持つモジュール。
storage_dtype (torch.dtype) — ストレージ用にモジュールをキャストするdtype（フォワードパスの前/後）。
compute_dtype (torch.dtype) — フォワードパスでの計算時にモジュールをキャストするdtype。
skip_modules_pattern (tuple[str, ...]、デフォルトはNone) — レイヤーワイズキャスティング処理中にスキップするモジュールの名前と一致させるためのパターンリストです。Noneに設定され、かつskip_modules_classesもNoneの場合、レイヤーワイズキャスティングはモジュールの内部サブモジュールではなく、モジュール自体に直接適用されます。
skip_modules_classes (tuple[type[torch.nn.Module], ...]、デフォルトはNone) — レイヤーワイズキャスト処理中にスキップするモジュールクラスのリスト。
non_blocking (bool、デフォルトはFalse) — Trueの場合、重みのキャスト操作はブロッキングなしで実行されます。

指定されたモジュールにレイヤー単位のキャストを適用します。ここで想定されるモジュールは PyTorch の nn.Module です。モデルを完全に量子化せずにメモリ要件を削減したい場合に有用です。モデルのパラメータを、例えば torch.float8_e4m3fn で保持し、フォワードパス中に torch.bfloat16 などの高精度にアップキャストし、メモリ削減を実現するために torch.float8_e4m3fn にダウキャストし直すことができます。

例:

>>> from accelerate.hooks import attach_layerwise_casting_hooks
>>> from transformers import AutoModelForCausalLM
>>> import torch

>>> # Model
>>> checkpoint = "EleutherAI/gpt-j-6B"
>>> model = AutoModelForCausalLM.from_pretrained(checkpoint)

>>> # Attach hooks and perform inference
>>> attach_layerwise_casting_hooks(model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16)
>>> with torch.no_grad():
...     model(...)

ユーザーは、ダウンキャストを回避したいモジュールを渡すこともできます。

>>> attach_layerwise_casting_hooks(
...     model, storage_dtype=torch.float8_e4m3fn, compute_dtype=torch.bfloat16, skip_modules_pattern=["norm"]
... )

フックの削除

remove_hook_from_module

accelerate.hooks.remove_hook_from_module

< source >

( module: Modulerecurse = False ) → torch.nn.Module

パラメータ

module (torch.nn.Module) — フックをアタッチするモジュール。
recurse (bool, 省略可) — フックを再帰的に削除するかどうか

戻り値

torch.nn.Module

同じモジュールで、フックが切り離された状態です（モジュールはインプレースで変更されるため、結果は破棄できます）。

add_hook_to_moduleによってモジュールにアタッチされたフックをすべて削除します。

remove_hook_from_submodules

accelerate.hooks.remove_hook_from_submodules

< source >

( module: Module )

パラメータ

module (torch.nn.Module) — すべてのフックを削除する対象のモジュール。

与えられたモデルのサブモジュールにアタッチされているすべてのフックを再帰的に削除します。

ユーティリティ

has_offloaded_params

accelerate.utils.has_offloaded_params

< source >

( module: Module ) → bool

パラメータ

module (torch.nn.Module) — オフロードフックのチェック対象となるモジュール。

戻り値

bool

モジュールにオフロードフックがあり、オフロードが有効な場合はTrue、それ以外の場合はFalse。

AlignDevicesHook がオフロード有効でアタッチされているかどうかを確認することにより、指定されたモジュールにオフロードされたパラメータがあるかどうかをチェックします

align_module_device

accelerate.utils.align_module_device

< source >

( module: Moduleexecution_device: torch.device | None = None )

パラメータ

module (torch.nn.Module) — アラインメント用のパラメータを持つモジュール。
execution_device (torch.device, 省略可) — 指定された場合、コンテキスト内のモジュールの実行デバイスを上書きします。指定されない場合、フックの実行デバイスを使用するか、そのまま通過します。

モジュールのパラメータを指定された実行デバイスに移動するコンテキストマネージャーです。

Update on GitHub

←ログ出力 Kwargsハンドラー→