目錄
  • 文件 >
  • Table Batched Embedding (TBE) 訓練模組
捷徑

Table Batched Embedding (TBE) 訓練模組

穩定版 API

class fbgemm_gpu.split_table_batched_embeddings_ops_training.SplitTableBatchedEmbeddingBagsCodegen(embedding_specs: List[Tuple[int, int, EmbeddingLocation, ComputeDevice]], feature_table_map: List[int] | None = None, cache_algorithm: CacheAlgorithm = CacheAlgorithm.LRU, cache_load_factor: float = 0.2, cache_sets: int = 0, cache_reserved_memory: float = 0.0, cache_precision: SparseType | None = None, weights_precision: SparseType = SparseType.FP32, output_dtype: SparseType = SparseType.FP32, enforce_hbm: bool = False, optimizer: EmbOptimType = EmbOptimType.EXACT_SGD, record_cache_metrics: RecordCacheMetrics | None = None, gather_uvm_cache_stats: bool | None = False, stochastic_rounding: bool = True, gradient_clipping: bool = False, max_gradient: float = 1.0, max_norm: float = 0.0, learning_rate: float = 0.01, eps: float = 1e-08, momentum: float = 0.9, weight_decay: float = 0.0, weight_decay_mode: WeightDecayMode = WeightDecayMode.NONE, eta: float = 0.001, beta1: float = 0.9, beta2: float = 0.999, ensemble_mode: EnsembleModeDefinition | None = None, emainplace_mode: EmainplaceModeDefinition | None = None, counter_based_regularization: CounterBasedRegularizationDefinition | None = None, cowclip_regularization: CowClipDefinition | None = None, pooling_mode: PoolingMode = PoolingMode.SUM, device: str | int | device | None = None, bounds_check_mode: BoundsCheckMode = BoundsCheckMode.WARNING, uvm_non_rowwise_momentum: bool = False, use_experimental_tbe: bool = False, prefetch_pipeline: bool = False, stats_reporter_config: TBEStatsReporterConfig | None = None, table_names: List[str] | None = None, optimizer_state_dtypes: Dict[str, SparseType] | None = None, multipass_prefetch_config: MultiPassPrefetchConfig | None = None, global_weight_decay: GlobalWeightDecayDefinition | None = None, uvm_host_mapped: bool = False, extra_optimizer_config: UserEnabledConfigDefinition | None = None, tbe_input_multiplexer_config: TBEInputMultiplexerConfig | None = None)[source]

Table Batched Embedding (TBE) 運算子。查詢一個或多個嵌入表。此模組適用於訓練。反向運算子與最佳化器融合。因此，嵌入表會在反向傳播期間更新。

參數:

  • embedding_specs (List[Tuple[int, int, EmbeddingLocation, ComputeDevice]]) –

    嵌入規格的列表。每個規格描述一個實體嵌入表的規格。每一個規格都是一個元組，包含嵌入列數、嵌入維度（必須是 4 的倍數）、表格放置位置 (EmbeddingLocation) 和運算裝置 (ComputeDevice)。

    可用的 EmbeddingLocation 選項如下

    1. DEVICE = 將嵌入表放置在 GPU 全域記憶體 (HBM) 中

    2. MANAGED = 將嵌入表放置在統一虛擬記憶體中（可從 GPU 和 CPU 存取）

    3. MANAGED_CACHING = 將嵌入表放置在統一虛擬記憶體中，並使用 GPU 全域記憶體 (HBM) 作為快取

    4. HOST = 將嵌入表放置在 CPU 記憶體 (DRAM) 中

    5. MTIA = 將嵌入表放置在 MTIA 記憶體中

    可用的 ComputeDevice 選項如下

    1. CPU = 在 CPU 上執行表格查詢

    2. CUDA = 在 GPU 上執行表格查詢

    3. MTIA = 在 MTIA 上執行表格查詢

  • feature_table_map (Optional[List[int]] = None) – 一個可選列表，用於指定特徵到表格的映射。 feature_table_map[i] 表示特徵 i 映射到的實體嵌入表。

  • cache_algorithm (CacheAlgorithm = CacheAlgorithm.LRU) –

    快取演算法（當 EmbeddingLocation 設定為 MANAGED_CACHING 時使用）。選項如下

    1. LRU = 最近最少使用

    2. LFU = 最不常使用

  • cache_load_factor (float = 0.2) – 當使用 EmbeddingLocation.MANAGED_CACHING 時，用於決定快取容量的因子。快取容量為 cache_load_factor * 所有嵌入表中的總列數

  • cache_sets (int = 0) – 快取組數（當 EmbeddingLocation 設定為 MANAGED_CACHING 時使用）

  • cache_reserved_memory (float = 0.0) – 在 HBM 中為非快取目的保留的記憶體量（當 EmbeddingLocation 設定為 MANAGED_CACHING 時使用）。

  • cache_precision (SparseType = SparseType.FP32) – 快取的資料類型（當 EmbeddingLocation 設定為 MANAGED_CACHING 時使用）。選項為 SparseType.FP32 和 SparseType.FP16

  • weights_precision (SparseType = SparseType.FP32) – 嵌入表的資料類型（也稱為權重）。選項為 SparseType.FP32 和 SparseType.FP16

  • output_dtype (SparseType = SparseType.FP32) – 輸出張量的資料類型。選項為 SparseType.FP32 和 SparseType.FP16

  • enforce_hbm (bool = False) – 若為 True，當使用 EmbeddingLocation.MANAGED_CACHING 時，將所有權重/動量放置在 HBM 中

  • optimizer (OptimType = OptimType.EXACT_SGD) –

    用於反向傳播中嵌入表更新的最佳化器。可用的 OptimType 選項如下

    1. ADAM = Adam

    2. EXACT_ADAGRAD = Adagrad

    3. EXACT_ROWWISE_ADAGRAD = Rowwise-Adagrad

    4. EXACT_SGD = SGD

    5. LAMB = Lamb

    6. LARS_SGD = LARS-SGD

    7. PARTIAL_ROWWISE_ADAM = Partial rowwise-Adam

    8. PARTIAL_ROWWISE_LAMB = Partial rowwise-Lamb

    9. ENSEMBLE_ROWWISE_ADAGRAD = Ensemble rowwise-Adagrad

    10. EMAINPLACE_ROWWISE_ADAGRAD = Ema inplace rowwise-Adagrad

    11. NONE = 在反向傳播中不應用最佳化器更新

    並輸出稀疏權重梯度

  • record_cache_metrics (Optional[RecordCacheMetrics] = None) – 如果 RecordCacheMetrics.record_cache_miss_counter 為 True，則記錄點擊次數、請求次數等，如果 RecordCacheMetrics.record_tablewise_cache_miss is True，則記錄類似的表格式指標

  • gather_uvm_cache_stats (Optional[bool] = False) – 若為 True，當 EmbeddingLocation 設定為 MANAGED_CACHING 時，收集快取統計資訊

  • stochastic_rounding (bool = True) – 若為 True，則對非 SparseType.FP32 的權重類型應用隨機捨入

  • gradient_clipping (bool = False) – 若為 True，則應用梯度裁剪

  • max_gradient (float = 1.0) – 梯度裁剪的值

  • max_norm (float = 0.0) – 最大範數值

  • learning_rate (float = 0.01) – 學習率

  • eps (float = 1.0e-8) – Adagrad、LAMB 和 Adam 使用的 epsilon 值。請注意，預設值與 torch.nn.optim.Adagrad 的預設值 1e-10 不同

  • momentum (float = 0.9) – LARS-SGD 使用的動量

  • weight_decay (float = 0.0) –

    LARS-SGD、LAMB、ADAM 和 rowwise-Adagrad 使用的權重衰減。

    1. EXACT_ADAGRAD、SGD、EXACT_SGD 不支援權重衰減

    2. LAMB、ADAM、PARTIAL_ROWWISE_ADAM、PARTIAL_ROWWISE_LAMB、LARS_SGD 支援解耦權重衰減

    3. EXACT_ROWWISE_ADAGRAD 同時支援 L2 和解耦權重衰減（透過 weight_decay_mode）

  • weight_decay_mode (WeightDecayMode = WeightDecayMode.NONE) – 權重衰減模式。選項為 WeightDecayMode.NONE、WeightDecayMode.L2 和 WeightDecayMode.DECOUPLE

  • eta (float = 0.001) – LARS-SGD 使用的 eta 值

  • beta1 (float = 0.9) – LAMB 和 ADAM 使用的 beta1 值

  • beta2 (float = 0.999) – LAMB 和 ADAM 使用的 beta2 值

  • ensemble_mode (Optional[EnsembleModeDefinition] = None) – 由 Ensemble Rowwise Adagrad 使用

  • emainplace_mode (Optional[EmainplaceModeDefinition] = None) – 由 EMA in-place Rowwise Adagrad 使用

  • counter_based_regularization (Optional[CounterBasedRegularizationDefinition] = None) – 由 Rowwise Adagrad 使用

  • cowclip_regularization (Optional[CowClipDefinition] = None) – 由 Rowwise Adagrad 使用

  • pooling_mode (PoolingMode = PoolingMode.SUM) –

    池化模式。可用的 PoolingMode 選項如下

    1. SUM = 總和池化

    2. MEAN = 平均池化

    3. NONE = 無池化（序列嵌入）

  • device (Optional[Union[str, int, torch.device]] = None) – 將張量放置於其上的目前裝置

  • bounds_check_mode (BoundsCheckMode = BoundsCheckMode.WARNING) –

    輸入檢查模式。可用的 BoundsCheckMode 選項如下

    1. NONE = 跳過邊界檢查

    2. FATAL = 當遇到無效的索引/偏移量時，拋出錯誤

    3. WARNING = 當遇到無效的索引/偏移量時，印出警告訊息並修正它（將無效的索引設定為零，並調整無效的偏移量使其在邊界內）

    4. IGNORE = 靜默修正無效的索引/偏移量（將無效的索引設定為零，並調整無效的偏移量使其在邊界內）

  • uvm_non_rowwise_momentum (bool = False) – 若為 True，將非 rowwise 動量放置在統一虛擬記憶體上

  • use_experimental_tbe (bool = False) – 若為 True，使用最佳化的 TBE 實作 (TBE v2)。請注意，這僅在 NVIDIA GPU 上支援。

  • prefetch_pipeline (bool = False) – 若為 True，當使用 EmbeddingLocation.MANAGED_CACHING 時，啟用快取預取管線。目前僅支援 LRU 快取策略。如果預取使用單獨的流，則必須設定預取函數的可選 forward_stream 引數。

  • stats_reporter_config (Optional[TBEStatsReporterConfig] = None) – TBE 統計資訊報告器的設定

  • table_names (Optional[List[str]] = None) – 此 TBE 中嵌入表名稱的列表

  • optimizer_state_dtypes (Optional[Dict[str, SparseType]] = None) – 最佳化器狀態資料類型字典。鍵是最佳化器狀態名稱，值是其對應的類型

  • multipass_prefetch_config (Optional[MultiPassPrefetchConfig] = None) – 多次傳遞快取預取設定（當使用 EmbeddingLocation.MANAGED_CACHING 時）

  • global_weight_decay (Optional[GlobalWeightDecayDefinition] = None) – 全域權重衰減的設定

  • uvm_host_mapped (bool = False) – 若為 True，則使用 malloc + cudaHostRegister 分配每個 UVM 張量。否則使用 cudaMallocManaged

  • None) (extra_optimizer_config Optional[UserEnabledConfigDefinition] =) –

    用於啟用最佳化器特定模式的額外設定。這些模式預設未啟用。- use_rowwise_bias_correction 在 Adam 中用於啟用 rowwise

    偏差校正計算

forward(indices: Tensor, offsets: Tensor, per_sample_weights: Tensor | None = None, feature_requires_grad: Tensor | None = None, batch_size_per_feature_per_rank: List[List[int]] | None = None, total_unique_indices: int | None = None) → Tensor[source]

執行前向傳遞函數，其功能為：

  1. 執行輸入邊界檢查

  2. 產生必要的變動批次大小嵌入 (VBE) metadata（如果使用 VBE）

  3. 將資料從 UVM 預先提取到快取（如果使用 EmbeddingLocation.MANAGED_CACHING 且使用者未明確預先提取資料）

  4. 透過調用對應的 Autograd 函數（基於選擇的優化器）來執行嵌入表查找

參數:

  • indices (Tensor) – 一個 1 維張量，包含要從所有嵌入表中查找的索引

  • offsets (Tensor) – 一個 1 維張量，包含索引的偏移量。形狀 (B * T + 1)，其中 B = 批次大小，T = 特徵數量。offsets[t * B + b + 1] - offsets[t * B + b] 是特徵 t 的包 b 的長度

  • per_sample_weights (Optional[Tensor]) – 一個可選的 1 維浮點張量，包含每個樣本的權重。如果為 None，將執行未加權嵌入查找。否則，將使用加權查找。此張量的長度必須與 indices 張量的長度相同。per_sample_weights[i] 的值將用於乘以查找行 indices[i] 中的每個元素，其中 0 <= i < len(per_sample_weights)。

  • feature_requires_grad (Optional[Tensor]) – 一個可選的 1 維張量，用於指示 per_sample_weights 是否需要梯度。張量的長度必須等於特徵的數量

  • batch_size_per_feature_per_rank (Optional[List[List[int]]]) – 一個可選的 2 維張量，包含每個 rank 和每個特徵的批次大小。如果為 None，TBE 會假設每個特徵都具有相同的批次大小，並從 offsets 形狀計算批次大小。否則，TBE 會假設不同的特徵可以有不同的批次大小，並使用變動批次大小嵌入查找模式 (VBE)。形狀為（特徵數量，rank 數量）。batch_size_per_feature_per_rank[f][r] 代表特徵 f 和 rank r 的批次大小

  • total_unique_indices (Optional[int]) – 一個可選的整數，表示唯一索引的總數。當使用 OptimType.NONE 時，必須設定此值。這是因為 TBE 需要此資訊才能在反向傳遞中分配權重梯度張量。

返回:

一個包含查找資料的 2 維張量。形狀 (B, total_D)，其中 B = 批次大小，total_D = 表中所有嵌入維度的總和

範例

>>> import torch
>>>
>>> from fbgemm_gpu.split_table_batched_embeddings_ops_common import (
>>>    EmbeddingLocation,
>>> )
>>> from fbgemm_gpu.split_table_batched_embeddings_ops_training import (
>>>    SplitTableBatchedEmbeddingBagsCodegen,
>>>    ComputeDevice,
>>> )
>>>
>>> # Two tables
>>> embedding_specs = [
>>>     (3, 8, EmbeddingLocation.DEVICE, ComputeDevice.CUDA),
>>>     (5, 4, EmbeddingLocation.MANAGED, ComputeDevice.CUDA)
>>> ]
>>>
>>> tbe = SplitTableBatchedEmbeddingBagsCodegen(embedding_specs)
>>> tbe.init_embedding_weights_uniform(-1, 1)
>>>
>>> print(tbe.split_embedding_weights())
[tensor([[-0.9426,  0.7046,  0.4214, -0.0419,  0.1331, -0.7856, -0.8124, -0.2021],
        [-0.5771,  0.5911, -0.7792, -0.1068, -0.6203,  0.4813, -0.1677,  0.4790],
        [-0.5587, -0.0941,  0.5754,  0.3475, -0.8952, -0.1964,  0.0810, -0.4174]],
       device='cuda:0'), tensor([[-0.2513, -0.4039, -0.3775,  0.3273],
        [-0.5399, -0.0229, -0.1455, -0.8770],
        [-0.9520,  0.4593, -0.7169,  0.6307],
        [-0.1765,  0.8757,  0.8614,  0.2051],
        [-0.0603, -0.9980, -0.7958, -0.5826]], device='cuda:0')]

>>> # Batch size = 3
>>> indices = torch.tensor([0, 1, 2, 0, 1, 2, 0, 3, 1, 4, 2, 0, 0],
>>>                        device="cuda",
>>>                        dtype=torch.long)
>>> offsets = torch.tensor([0, 2, 5, 7, 9, 12, 13],
>>>                        device="cuda",
>>>                        dtype=torch.long)
>>>
>>> output = tbe(indices, offsets)
>>>
>>> # Batch size = 3, total embedding dimension = 12
>>> print(output.shape)
torch.Size([3, 12])

>>> print(output)
tensor([[-1.5197,  1.2957, -0.3578, -0.1487, -0.4873, -0.3044, -0.9801,  0.2769,
         -0.7164,  0.8528,  0.7159, -0.6719],
        [-2.0784,  1.2016,  0.2176,  0.1988, -1.3825, -0.5008, -0.8991, -0.1405,
         -1.2637, -0.9427, -1.8902,  0.3754],
        [-1.5013,  0.6105,  0.9968,  0.3057, -0.7621, -0.9821, -0.7314, -0.6195,
         -0.2513, -0.4039, -0.3775,  0.3273]], device='cuda:0',
       grad_fn=<CppNode<SplitLookupFunction_sgd_Op>>)
set_learning_rate(lr: float) → None[source]

設定學習率。

參數:

lr (float) – 要設定的學習率值

set_optimizer_step(step: int) → None[source]

設定優化器步驟。

參數:

step (int) – 要設定的步驟值

split_embedding_weights() → List[Tensor][source]

傳回嵌入權重（view）的列表，依表格分割

返回:

權重列表。長度 = 表格數量

split_optimizer_states() → List[List[Tensor]][source]

傳回優化器狀態（view）的列表，依表格分割

返回:

狀態列表的列表。形狀 = (表格數量，狀態數量)。

以下顯示每個優化器的狀態列表（依傳回順序排列）

  1. ADAM: momentum1, momentum2

  2. EXACT_ADAGRAD: momentum1

  3. EXACT_ROWWISE_ADAGRAD: momentum1 (逐行), prev_iter (逐行；僅當使用 WeightDecayMode = COUNTER 或 COWCLIP 或 global_weight_decay 不為 None 時), row_counter (逐行；僅當使用 WeightDecayMode = COUNTER 或 COWCLIP 時)

  4. EXACT_SGD: 無狀態

  5. LAMB: momentum1, momentum2

  6. LARS_SGD: momentum1

  7. PARTIAL_ROWWISE_ADAM: momentum1, momentum2 (逐行)

  8. PARTIAL_ROWWISE_LAMB: momentum1, momentum2 (逐行)

  9. ENSEMBLE_ROWWISE_ADAGRAD: momentum1 (逐行), momentum2

  10. NONE: 無狀態 (拋出錯誤)

update_hyper_parameters(params_dict: Dict[str, float]) → None[source]

從外部控制流程設定超參數。

參數:

params_dict (Dict[str, float]) – 包含超參數名稱及其值的字典

其他 API

文件

存取 PyTorch 的完整開發人員文件

查看文件

教學課程

取得針對初學者和進階開發人員的深入教學課程

查看教學課程

資源

尋找開發資源並獲得問題解答

查看資源

為了分析流量並優化您的體驗，我們在這個網站上使用 Cookie。透過點擊或瀏覽，您同意允許我們使用 Cookie。作為本網站目前的維護者，Facebook 的 Cookie 政策適用。了解更多資訊，包括可用的控制項：Cookie 政策。