[feat] support ftrl sparse optimizer for dynamicemb

docs/source/feature/dynamicemb.md

@@ -9,10 +9,14 @@ cu126/cu129/cu130 镜像已预装 dynamicemb，其它环境需先安装如下whl
 pip install dynamicemb==0.1.0+20260920.9643985.${DEVICE} -f https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/${DEVICE}/repo.html
 ```
 
+支持配置 DynamicEmbedding 的特征类型为 `id_feature`、`combo_feature`、`lookup_feature`、`match_feature`、`regex_replace_feature`、`custom_feature`、`bool_mask_feature`，其余特征类型（如配置了 `boundaries` 的 `raw_feature`、`expr_feature`、`tokenize_feature`、`combine_feature`）暂不支持。
+
 注：同一个 FeatureGroup 中若存在多个配置了 DynamicEmbedding 的特征，底层 dynamicemb 会自动将这些表融合到同一份存储里（table fusion），共享 cache/admission counter，降低显存占用并减少内存碎片，无需额外配置。
 
 注：配置了 DynamicEmbedding 的模型导出时需设置环境变量 `USE_DISTRIBUTED_EMBEDDING=1`，使用分布式 embedding 导出模式，详见[模型导出](../usage/export.md)的环境变量章节。
 
+注：DynamicEmbedding 表额外支持 `ftrl_optimizer`（普通 Embedding 表不支持），配置方式见[优化器](../models/optimizer.md)文档。
+
 以id_feature的配置为例，DynamicEmbedding 只需在id_feature新增一个dynamicemb的配置字段
 
 ```

docs/source/models/optimizer.md

@@ -52,7 +52,9 @@ train_config {
 
   **Note**: 被分片为`data_parallel`的Embedding表不受sparse_optimizer管理，实际由dense_optimizer更新，并跟随dense_optimizer的LR策略，详见[训练文档](../usage/train.md)的Embedding分片约束章节
 
-  **Note**: `adagrad_optimizer`和`rowwise_adagrad_optimizer`的`initial_accumulator_value`对齐TensorFlow Adagrad的同名参数（TF默认0.1，TorchEasyRec默认0.0），对普通Embedding表和[dynamicemb](../feature/dynamicemb.md)表同时生效：普通Embedding表在建表时把整个accumulator初始化为该值，dynamicemb表则在key首次写入时把该key的accumulator初始化为该值
+  **Note**: `adagrad_optimizer`和`rowwise_adagrad_optimizer`的`initial_accumulator_value`对齐TensorFlow Adagrad的同名参数（TF默认0.1，TorchEasyRec默认0.0），对普通Embedding表和[dynamicemb](../feature/dynamicemb.md)表同时生效：普通Embedding表在建表时把整个accumulator初始化为该值，dynamicemb表则在key首次写入时把该key的accumulator初始化为该值。`ftrl_optimizer`也用该字段初始化它的accumulator
+
+  **Note**: `ftrl_optimizer`（FTRL-Proximal，McMahan et al. 2013）**只支持[dynamicemb](../feature/dynamicemb.md)表**，FBGEMM没有FTRL的embedding kernel，模型中只要还有一张非dynamicemb的sparse表，训练会在plan阶段直接报错并给出表名。可配置`dynamicemb`的特征类型见[dynamicemb文档](../feature/dynamicemb.md)，配置了`boundaries`的`raw_feature`等不支持dynamicemb的特征，无法与`ftrl_optimizer`一起使用。被分片为`data_parallel`的表不受此限制（由dense_optimizer更新）
 
 - dense_optimizer
 

tzrec/optim/optimizer.py

@@ -9,7 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import logging
-from typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Iterable, Optional, Union
 
 import torch
 from fbgemm_gpu import split_table_batched_embeddings_ops_training
@@ -19,6 +19,8 @@
 )
 from torch import nn
 from torch.amp import GradScaler
+from torch.optim.optimizer import Optimizer
+from torchrec.distributed.utils import _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE
 from torchrec.optim import KeyedOptimizer, OptimizerWrapper
 
 
@@ -67,6 +69,51 @@ def step(self, closure: Any = None) -> None:
                 self._optimizer.step(closure=closure)
 
 
+class FTRL(Optimizer):
+    """Placeholder for the dynamicemb FTRL sparse embedding optimizer.
+
+    FBGEMM has no FTRL kernel, so torchrec ships no FTRL wrapper to reuse. Like
+    torchrec's own placeholders this class never runs: it names the optimizer so
+    that the sharding plan can resolve it, and the update happens inside the
+    dynamicemb table.
+
+    Args:
+        params (Iterable[nn.Parameter]): parameters to attach the optimizer to.
+        **kwargs: fused params, forwarded to the dynamicemb table.
+    """
+
+    def __init__(self, params: Iterable[nn.Parameter], **kwargs: Any) -> None:
+        self._params = params
+        self._kwargs = kwargs
+
+    # pyrefly: ignore[bad-override]  # matches torchrec's placeholder optimizers
+    def step(self, closure: Optional[Callable[[], float]] = None) -> None:
+        """Step, never reached, the dynamicemb table applies the update."""
+        raise NotImplementedError
+
+
+def register_ftrl_emb_opt_type() -> None:
+    """Map :class:`FTRL` to dynamicemb's optimizer type in torchrec's table.
+
+    torchrec derives an embedding table's fused `optimizer` param from the
+    in-backward optimizer class, and does so unconditionally on the
+    EmbeddingCollection path, so FTRL has to be registered in that table rather
+    than injected into the fused params. Must be called before planning.
+
+    Raises:
+        RuntimeError: dynamicemb is missing or predates its FTRL support.
+    """
+    try:
+        from dynamicemb import DynamicEmbOptimType
+    except ImportError as e:
+        raise RuntimeError(
+            "sparse ftrl_optimizer requires dynamicemb >= "
+            "0.1.0+20260920.9643985; FBGEMM has no FTRL embedding kernel. "
+            "Please reinstall dynamicemb, see docs/source/feature/dynamicemb.md."
+        ) from e
+    _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE[FTRL] = DynamicEmbOptimType.FTRL
+
+
 # The Adagrad optimizer in TensorFlow includes the parameter
 # `initial_accumulator_value`, with a default value of 0.1.
 # Here, we patch the fbgemm embedding optimizer state split helper
@@ -75,11 +122,13 @@ def step(self, closure: Any = None) -> None:
 
 
 def set_sparse_init_accumulator_value(value: float) -> None:
-    """Record the Adagrad accumulator initial value for embedding tables built later.
+    """Record the accumulator initial value for embedding tables built later.
 
     Takes effect at table build time, in the ``apply_split_helper`` patch below and
     in ``dynamicemb_util``'s plan-time fused params, so it must be set before
-    planning; FBGEMM TBE has no such kwarg, hence this module-level switch.
+    planning; FBGEMM TBE has no such kwarg, hence this module-level switch. Used
+    by Adagrad and, on dynamicemb tables, by FTRL for its squared-gradient
+    accumulator.
 
     Args:
         value: accumulator initial value, 0.0 for optimizers without one.
@@ -89,7 +138,7 @@ def set_sparse_init_accumulator_value(value: float) -> None:
 
 
 def sparse_init_accumulator_value() -> float:
-    """Sparse Adagrad accumulator initial value, 0.0 when not configured."""
+    """Sparse accumulator initial value, 0.0 when not configured."""
     return _sparse_init_accumulator_value
 
 

tzrec/optim/optimizer_builder.py

@@ -21,7 +21,11 @@
 from torchrec.optim.keyed import KeyedOptimizerWrapper
 
 from tzrec.optim.lr_scheduler import BaseLR
-from tzrec.optim.optimizer import set_sparse_init_accumulator_value
+from tzrec.optim.optimizer import (
+    FTRL,
+    register_ftrl_emb_opt_type,
+    set_sparse_init_accumulator_value,
+)
 from tzrec.protos import optimizer_pb2
 from tzrec.utils.config_util import config_to_kwargs
 from tzrec.utils.logging_util import logger
@@ -92,6 +96,9 @@ def create_sparse_optimizer(
         # FBGEMM reuses the beta1 OptimizerArgs slot for RMSProp's alpha.
         optimizer_kwargs["beta1"] = optimizer_kwargs.pop("alpha")
         return optimizers.RMSProp, optimizer_kwargs
+    elif optimizer_type == "ftrl_optimizer":
+        register_ftrl_emb_opt_type()
+        return FTRL, optimizer_kwargs
     else:
         raise ValueError(f"Unknown optimizer: {optimizer_type}")
 

tzrec/optim/optimizer_builder_test.py

@@ -15,20 +15,32 @@
 from parameterized import param, parameterized
 from torch import Tensor, nn
 from torch.optim import Optimizer
+from torchrec.distributed.utils import _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE
 from torchrec.optim.keyed import KeyedOptimizerWrapper
 
 from tzrec.optim import optimizer_builder
 from tzrec.optim.optimizer import (
+    FTRL,
     set_sparse_init_accumulator_value,
     sparse_init_accumulator_value,
 )
 from tzrec.protos import optimizer_pb2
-from tzrec.utils.test_util import parameterized_name_func
+from tzrec.utils.test_util import mark_ci_scope, parameterized_name_func
+
+try:
+    from dynamicemb import DynamicEmbOptimType
+
+    has_dynamicemb = True
+except ImportError:
+    has_dynamicemb = False
 
 
 class OpimizerBuilderTest(unittest.TestCase):
     def tearDown(self):
         set_sparse_init_accumulator_value(0.0)
+        # The whole suite runs in one process, so leaving FTRL registered would
+        # mask a missing register_ftrl_emb_opt_type() call in a later test.
+        _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE.pop(FTRL, None)
 
     @parameterized.expand(
         [
@@ -73,6 +85,40 @@ def test_create_sparse_optimizer_init_accumulator_value(
         self.assertNotIn("initial_accumulator_value", kwargs)
         self.assertAlmostEqual(sparse_init_accumulator_value(), expected)
 
+    @unittest.skipUnless(
+        has_dynamicemb, "dynamicemb with FTRL support is not installed."
+    )
+    @mark_ci_scope("gpu")
+    def test_create_sparse_optimizer_ftrl(self):
+        from torchrec.distributed.utils import optimizer_type_to_emb_opt_type
+
+        optimizer_config = optimizer_pb2.SparseOptimizer(
+            ftrl_optimizer=optimizer_pb2.FusedFTRLOptimizer(
+                lr=0.01,
+                learning_rate_power=-0.4,
+                ftrl_beta=1.0,
+                l1_reg=0.01,
+                l2_reg=0.02,
+                initial_accumulator_value=0.1,
+            ),
+            constant_learning_rate=optimizer_pb2.ConstantLR(),
+        )
+        optim_cls, kwargs = optimizer_builder.create_sparse_optimizer(optimizer_config)
+        self.assertIs(optim_cls, FTRL)
+        # torchrec turns the class into the fused optimizer param of the table.
+        self.assertIs(
+            optimizer_type_to_emb_opt_type(optim_cls), DynamicEmbOptimType.FTRL
+        )
+        self.assertNotIn("initial_accumulator_value", kwargs)
+        self.assertAlmostEqual(sparse_init_accumulator_value(), 0.1)
+        self.assertAlmostEqual(kwargs["lr"], 0.01)
+        self.assertAlmostEqual(kwargs["learning_rate_power"], -0.4)
+        self.assertAlmostEqual(kwargs["ftrl_beta"], 1.0)
+        self.assertAlmostEqual(kwargs["l1_reg"], 0.01)
+        self.assertAlmostEqual(kwargs["l2_reg"], 0.02)
+        # apply_optimizer_in_backward constructs the class with these kwargs.
+        optim_cls([nn.Parameter(torch.zeros(4, 8))], **kwargs)
+
     def test_create_part_optimizer(self):
         pattern1 = "model.dbmtl.task(.*)"
         pattern2 = "model.dbmtl.mmoe(.*)"

tzrec/protos/optimizer.proto

@@ -13,6 +13,7 @@ message SparseOptimizer {
         FusedRowWiseAdagradOptimizer rowwise_adagrad_optimizer = 8;
         FusedAdadeltaOptimizer adadelta_optimizer = 9;
         FusedRMSpropOptimizer rmsprop_optimizer = 10;
+        FusedFTRLOptimizer ftrl_optimizer = 11;
     }
     oneof learning_rate {
         ConstantLR constant_learning_rate = 101;
@@ -160,6 +161,26 @@ message FusedRMSpropOptimizer {
     optional float max_gradient = 6 [default = 1.0];
 }
 
+// FTRL-Proximal, Algorithm 1 of McMahan et al. 2013, with the paper's fixed
+// square root generalized to an arbitrary exponent. dynamicemb tables only.
+message FusedFTRLOptimizer {
+    // alpha in the paper, must be > 0 because the update divides by it.
+    required float lr = 1 [default = 0.002];
+    // exponent on the accumulator in the learning rate, must be <= 0.
+    // -0.5 recovers the paper's alpha / (beta + sqrt(n)).
+    optional float learning_rate_power = 2 [default = -0.5];
+    // beta in the paper, bounds the learning rate while the accumulator is small.
+    optional float ftrl_beta = 3 [default = 0.0];
+    // lambda1 in the paper, L1 regularization.
+    optional float l1_reg = 4 [default = 0.0];
+    // lambda2 in the paper, L2 regularization.
+    optional float l2_reg = 5 [default = 0.0];
+    optional bool gradient_clipping = 6 [default = false];
+    optional float max_gradient = 7 [default = 1.0];
+    // initial value of the squared-gradient accumulator.
+    optional float initial_accumulator_value = 8 [default = 0.0];
+}
+
 message SGDOptimizer {
     required float lr = 1 [default = 0.002];
     optional float momentum = 2 [default = 0.9];