Accelerate

( include_buffers: bool = None )

引數

include_buffers (bool, optional) — 是否在初始化時也將所有緩衝區放在元裝置上。

一個上下文管理器，在此之下，模型初始化時所有引數都位於元裝置上，從而建立一個空模型。當僅僅初始化模型就會耗盡可用 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)])

在此上下文管理器下建立的任何模型都沒有權重。因此，你不能對其執行類似 model.to(some_device) 的操作。要在空模型中載入權重，請參閱 load_checkpoint_and_dispatch()。請確保覆蓋 load_checkpoint_and_dispatch() 的預設 device_map 引數，否則不會呼叫排程。

cpu_offload

accelerate.cpu_offload

( model: Module execution_device: typing.Optional[torch.device] = None offload_buffers: bool = False state_dict: typing.Optional[dict[str, torch.Tensor]] = None preload_module_classes: typing.Optional[list[str]] = None )

引數

model (torch.nn.Module) — 要解除安裝的模型。
execution_device (torch.device, optional) — 模型前向傳播將在其上執行的裝置（應為 GPU）。將預設為模型的第一個引數所在的裝置。
offload_buffers (bool, optional, 預設為 False) — 是否將緩衝區與模型引數一起解除安裝。
state_dict (Dict[str, torch.Tensor], optional) — 將儲存在 CPU 上的模型的狀態字典。
preload_module_classes (List[str], optional) — 一個類名列表，這些類的例項應在前向傳播開始時載入其所有權重（即使在子模組中）。這應該只用於那些具有已註冊但在前向傳播期間未直接呼叫的子模組的類，例如，如果註冊了一個 `dense` 線性層，但在前向傳播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接呼叫 `dense`。

為一個模型啟用完全 CPU 解除安裝。因此，模型的所有引數將被解除安裝，並且只保留模型狀態字典的一個副本。在前向傳播期間，引數將從該狀態字典中提取，並根據需要放到指定的執行裝置上，然後再重新解除安裝。

cpu_offload_with_hook

accelerate.cpu_offload_with_hook

( model: Module execution_device: typing.Union[int, str, torch.device, NoneType] = None prev_module_hook: typing.Optional[accelerate.hooks.UserCpuOffloadHook] = None )

引數

model (torch.nn.Module) — 要解除安裝的模型。
execution_device(str, int 或 torch.device, optional) — 模型應在其上執行的裝置。如果 MPS 裝置可用，則預設為 MPS 裝置；否則，如果有 GPU，則為 GPU 0；最後為 CPU。
prev_module_hook (UserCpuOffloadHook, optional) — 此函式為正在執行的流水線中前一個模型返回的鉤子。如果傳遞了此引數，其 offload 方法將在附加此鉤子的模型的前向傳播之前被呼叫。

將模型解除安裝到 CPU 上，並在執行時將其放回執行裝置。與 cpu_offload() 的區別在於，模型在前向傳播後會保留在執行裝置上，只有在呼叫返回的 `hook` 的 `offload` 方法時才會再次解除安裝。這對於迴圈執行模型的流水線很有用。

示例

model_1, hook_1 = cpu_offload_with_hook(model_1, cuda_device)
model_2, hook_2 = cpu_offload_with_hook(model_2, cuda_device, prev_module_hook=hook_1)
model_3, hook_3 = cpu_offload_with_hook(model_3, cuda_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

( model: Module offload_dir: typing.Union[str, os.PathLike] execution_device: typing.Optional[torch.device] = None offload_buffers: bool = False preload_module_classes: typing.Optional[list[str]] = None )

引數

model (torch.nn.Module) — 要解除安裝的模型。
offload_dir (str 或 os.PathLike) — 用於解除安裝模型權重的資料夾（或者模型權重已經解除安裝到的位置）。
execution_device (torch.device, optional) — 模型前向傳播將在其上執行的裝置（應為 GPU）。將預設為模型的第一個引數所在的裝置。
offload_buffers (bool, optional, 預設為 False) — 是否將緩衝區與模型引數一起解除安裝。
preload_module_classes (List[str], optional) — 一個類名列表，這些類的例項應在前向傳播開始時載入其所有權重（即使在子模組中）。這應該只用於那些具有已註冊但在前向傳播期間未直接呼叫的子模組的類，例如，如果註冊了一個 `dense` 線性層，但在前向傳播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接呼叫 `dense`。

為模型啟用完全磁碟解除安裝。因此，模型的所有引數都將作為記憶體對映陣列解除安裝到給定資料夾中。在前向傳播期間，將從該資料夾訪問引數，並根據需要將其放置到指定的執行裝置上，然後再次解除安裝。

dispatch_model

accelerate.dispatch_model

( model: Module device_map: dict main_device: typing.Optional[torch.device] = None state_dict: typing.Optional[dict[str, torch.Tensor]] = None offload_dir: typing.Union[str, os.PathLike, NoneType] = None offload_index: typing.Optional[dict[str, str]] = None offload_buffers: bool = False skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None force_hooks: bool = False )

引數

model (torch.nn.Module) — 要排程的模型。
device_map (Dict[str, Union[str, int, torch.device]]) — 一個字典，將模型 `state_dict` 中的模組名稱對映到它們應該放置的裝置。請注意，即使 `"disk"` 不是 `torch.device` 的有效值，它也被接受。
main_device (str, int 或 torch.device, optional) — 主要執行裝置。將預設為 `device_map` 中第一個不同於 `"cpu"` 或 `"disk"` 的裝置。
state_dict (Dict[str, torch.Tensor], optional) — 將儲存在 CPU 上的模型部分的狀態字典。
offload_dir (str 或 os.PathLike) — 用於解除安裝模型權重的資料夾（或者模型權重已經解除安裝到的位置）。
offload_index (Dict, optional) — 一個從權重名稱到其資訊（`dtype`/`shape` 或 safetensors 檔名）的字典。將預設為 `save_folder` 中儲存的索引。
offload_buffers (bool, optional, 預設為 False) — 是否將緩衝區與模型引數一起解除安裝。
skip_keys (str 或 List[str], optional) — 在裝置之間移動輸入或輸出時要忽略的鍵的列表。
preload_module_classes (List[str], optional) — 一個類名列表，這些類的例項應在前向傳播開始時載入其所有權重（即使在子模組中）。這應該只用於那些具有已註冊但在前向傳播期間未直接呼叫的子模組的類，例如，如果註冊了一個 `dense` 線性層，但在前向傳播中，`dense.weight` 和 `dense.bias` 在某些操作中使用，而不是直接呼叫 `dense`。
force_hooks (bool, optional, 預設為 False) — 是否強制將裝置鉤子附加到模型，即使所有層都被排程到單個裝置。

根據給定的裝置對映排程模型。模型的層可能分佈在不同的 GPU 上，解除安裝到 CPU 甚至磁碟上。

load_checkpoint_and_dispatch

accelerate.load_checkpoint_and_dispatch

( model: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Union[str, dict[str, typing.Union[int, str, torch.device]], NoneType] = None max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None offload_buffers: bool = False dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: typing.Optional[bool] = None skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None force_hooks: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: bool = False )

引數

model (torch.nn.Module) — 我們希望載入檢查點的模型。
checkpoint (str 或 os.PathLike) — 要載入的資料夾檢查點。它可以是：
- 包含完整模型狀態字典的檔案的路徑
- 包含分片檢查點索引的 .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 狀態字典解除安裝到硬碟上，以避免在 CPU 狀態字典的權重加上最大分片的權重不適合時耗盡 CPU RAM。如果選擇的裝置對映包含 "disk" 值，則預設為 True。
skip_keys (str 或 List[str], 可選) — 在裝置之間移動輸入或輸出時要忽略的鍵的列表。
preload_module_classes (List[str], 可選) — 一個類的列表，其例項應在 forward 開始時載入所有權重（即使在子模組中）。這隻應用於那些已註冊但在 forward 期間未直接呼叫的子模組的類，例如，如果註冊了一個 dense 線性層，但在 forward 時，在某些操作中使用了 dense.weight 和 dense.bias，而不是直接呼叫 dense。
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 中的張量逐一廣播到其他 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

( model: Module checkpoint: typing.Union[str, os.PathLike] device_map: typing.Optional[dict[str, typing.Union[int, str, torch.device]]] = None offload_folder: typing.Union[str, os.PathLike, NoneType] = None dtype: typing.Union[str, torch.dtype, NoneType] = None offload_state_dict: bool = False offload_buffers: bool = False keep_in_fp32_modules: list = None offload_8bit_bnb: bool = False strict: bool = False full_state_dict: bool = True broadcast_from_rank0: bool = False )

引數

model (torch.nn.Module) — 我們希望載入檢查點的模型。
checkpoint (str 或 os.PathLike) — 要載入的資料夾檢查點。它可以是：
- 包含完整模型狀態字典的檔案的路徑
- 包含分片檢查點索引的 .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 狀態字典解除安裝到硬碟上，以避免在 CPU 狀態字典的權重加上最大分片的權重不適合時耗盡 CPU RAM。
offload_buffers (bool, 可選, 預設為 False) — 是否將緩衝區包含在解除安裝到磁碟的權重中。
keep_in_fp32_modules(List[str], 可選) — 保持 torch.float32 資料型別的模組列表。
offload_8bit_bnb (bool, 可選) — 是否在 cpu/磁碟上啟用 8-bit 模組的解除安裝。
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 中的張量逐一廣播到其他 rank。其他 rank 將接收張量並根據模型中的本地分片進行分片（如果適用）。

將一個（可能是分片的）檢查點載入到模型中，載入時可能會將權重發送到給定的裝置。

跨裝置載入後，您仍需要在模型上呼叫 dispatch_model() 才能使其執行。要將檢查點載入和分派組合在一次呼叫中，請使用 load_checkpoint_and_dispatch()。

infer_auto_device_map

accelerate.infer_auto_device_map

( model: Module max_memory: typing.Optional[dict[typing.Union[int, str], typing.Union[int, str]]] = None no_split_module_classes: typing.Optional[list[str]] = None dtype: typing.Union[str, torch.dtype, NoneType] = None special_dtypes: typing.Optional[dict[str, typing.Union[str, torch.dtype]]] = None verbose: bool = False clean_result: bool = True offload_buffers: bool = False fallback_allocation: bool = False )

引數

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）。
verbose (bool, 可選, 預設為 False) — 是否在函式構建 device_map 時提供除錯語句。
clean_result (bool, 可選, 預設為 True) — 透過將所有放置在同一裝置上的子模組分組來清理生成的 device_map。
offload_buffers (bool, 可選, 預設為 False) — 在解除安裝到 CPU 或硬碟的層中，是否也將緩衝區與引數一起解除安裝。
fallback_allocation (bool, 可選, 預設為 False) — 當常規分配失敗時，嘗試使用 BFS 分配一個適合大小限制的模組。

為給定的模型計算一個裝置對映，優先考慮 GPU，然後解除安裝到 CPU，最後解除安裝到磁碟，以便

我們不會超過任何 GPU 的可用記憶體。
如果需要解除安裝到 CPU，GPU 0 上始終留有空間來放回解除安裝到 CPU 的最大尺寸的層。
如果需要解除安裝到 CPU，我們不會超過 CPU 上可用的 RAM。
如果需要解除安裝到磁碟，CPU 上始終留有空間來放回解除安裝到磁碟的最大尺寸的層。

所有的計算都是透過分析模型引數的大小和資料型別來完成的。因此，模型可以位於元裝置上（就像在 init_empty_weights 上下文管理器中初始化時一樣）。

鉤子 (Hooks)

ModelHook

class accelerate.hooks.ModelHook

( )

一個包含在模型 forward 方法前後執行的回撥的鉤子。與 PyTorch 現有的鉤子不同的是，它們會傳遞 kwargs。

類屬性

no_grad (bool, 可選, 預設為 False) — 是否在 torch.no_grad() 上下文管理器下執行實際的前向傳播。

detach_hook

( module )

引數

module (torch.nn.Module) — 與此鉤子分離的模組。

當鉤子從模組上分離時執行。

init_hook

( module )

引數

module (torch.nn.Module) — 附加此鉤子的模組。

當鉤子附加到模組時執行。

post_forward

( module output ) → Any

引數

module (torch.nn.Module) — 剛剛執行完前向傳播的模組。
output (Any) — 模組的輸出。

Any

處理後的 output。

在模型的 forward 方法之後立即執行。

pre_forward

( 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

( execution_device: typing.Union[int, str, torch.device, NoneType] = None offload: bool = False io_same_device: bool = False weights_map: typing.Optional[collections.abc.Mapping] = None offload_buffers: bool = False place_submodules: bool = False skip_keys: typing.Union[str, list[str], NoneType] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None )

引數

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

( *hooks )

一個可以包含多個鉤子的鉤子，並在每個事件中遍歷它們。

LayerwiseCastingHook

class accelerate.hooks.LayerwiseCastingHook

( storage_dtype: dtype compute_dtype: dtype non_blocking: bool )

一個將模組權重轉換為高精度 dtype 進行計算，並轉換為低精度 dtype 進行儲存的鉤子。此過程可能會導致輸出質量下降，但可以顯著減少記憶體佔用。

新增鉤子

add_hook_to_module

accelerate.hooks.add_hook_to_module

( module: Module hook: ModelHook append: bool = False ) → torch.nn.Module

引數

module (torch.nn.Module) — 要附加鉤子的模組。
hook (ModelHook) — 要附加的鉤子。
append (bool, 可選, 預設為 False) — 鉤子是否應與現有鉤子（如果模組已包含鉤子）連結，或者不連結。

torch.nn.Module

帶有附加鉤子的相同模組（模組是原地修改的，因此結果可以丟棄）。

向給定模組新增一個鉤子。這將重寫模組的 forward 方法以包含該鉤子。要移除此行為並恢復原始的 forward 方法，請使用 remove_hook_from_module。

如果模組已包含一個鉤子，預設情況下會用傳遞的新鉤子替換它。要將兩個鉤子連結在一起，請傳遞 append=True，這樣它會將當前鉤子和新鉤子連結成一個 SequentialHook 類的例項。

attach_execution_device_hook

accelerate.hooks.attach_execution_device_hook

( module: Module execution_device: typing.Union[int, str, torch.device] skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None )

引數

module (torch.nn.Module) — 我們要附加鉤子的模組。
execution_device (int, str 或 torch.device) — 在前向傳播之前，輸入和模型權重應放置的裝置。
skip_keys (str 或 List[str], 可選) — 在裝置之間移動輸入或輸出時要忽略的鍵列表。
preload_module_classes (List[str], 可選) — 一個類列表，其類的例項應在前向傳播開始時載入所有權重（即使在子模組中）。這應該只用於那些具有已註冊但在前向傳播期間未直接呼叫的子模組的類，例如，如果註冊了一個 dense 線性層，但在前向傳播時，dense.weight 和 dense.bias 用於某些操作，而不是直接呼叫 dense。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 可選, 預設為 None) — 一個從資料指標到裝置字典的對映，該字典包含已排程的繫結權重。對於給定的執行裝置，此引數可用於為所有共享權重重用第一個可用的指標，而不是複製記憶體。

遞迴地將 AlignDevicesHook 附加到給定模型的所有子模組，以確保它們具有正確的執行裝置

attach_align_device_hook

accelerate.hooks.attach_align_device_hook

( module: Module execution_device: typing.Optional[torch.device] = None offload: bool = False weights_map: typing.Optional[collections.abc.Mapping] = None offload_buffers: bool = False module_name: str = '' skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = 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.weight 和 dense.bias 用於某些操作，而不是直接呼叫 dense。
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

( module: Module execution_device: typing.Union[torch.device, dict[str, torch.device], NoneType] = None offload: typing.Union[bool, dict[str, bool]] = False weights_map: Mapping = None offload_buffers: bool = False module_name: str = '' skip_keys: typing.Union[str, list[str], NoneType] = None preload_module_classes: typing.Optional[list[str]] = None tied_params_map: typing.Optional[dict[int, dict[torch.device, torch.Tensor]]] = None )

引數

module (torch.nn.Module) — 我們要附加鉤子的模組。
execution_device (torch.device 或 Dict[str, 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], 可選) — 一個類列表，其例項應在 forward 開始時載入其所有權重（即使在子模組中）。這應該只用於那些子模組已註冊但在 forward 期間未被直接呼叫的類，例如，如果一個 dense 線性層已註冊，但在 forward 期間，dense.weight 和 dense.bias 在某些操作中使用，而不是直接呼叫 dense。
tied_params_map (Optional[Dict[int, Dict[torch.device, torch.Tensor]]], 可選, 預設為 None) — 一個從資料指標到裝置字典的對映，該字典包含已分派的繫結權重。對於給定的執行裝置，此引數可用於為所有共享權重重用第一個可用的指標，而不是複製記憶體。

根據需要將 AlignDevicesHook 附加到給定模型的所有塊上。

attach_layerwise_casting_hooks

accelerate.big_modeling.attach_layerwise_casting_hooks