[feat] support ftrl sparse optimizer for dynamicemb

[tiankongdeguiji]

dynamicemb gained FTRL-Proximal in NVIDIA/recsys-examples#487. It lives in dynamicemb's own DynamicEmbOptimType rather than FBGEMM's EmbOptimType, because FBGEMM has no FTRL embedding kernel. This exposes it as train_config.sparse_optimizer.ftrl_optimizer. The wheel that carries it (0.1.0+20260920.9643985) is already pinned on master as of #680, so this PR no longer touches the requirements.

FTRL is a common choice for the sparse part of CTR models: it re-solves each weight from (linear, accum) every step instead of nudging it, which is what lets l1_reg pin coordinates to exactly zero.

train_config {
    sparse_optimizer {
        ftrl_optimizer {
            lr: 0.01
            learning_rate_power: -0.5
            ftrl_beta: 1.0
            l1_reg: 0.001
            l2_reg: 0.001
            initial_accumulator_value: 0.1
        }
        constant_learning_rate {
        }
    }
}

Notes on the approach

Why a placeholder class instead of naming the type in the optimizer kwargs. torchrec derives a table's fused optimizer param from the in-backward optimizer class. embeddingbag.py skips that when the caller already put an optimizer key in the kwargs, but embedding.py overwrites it unconditionally, so injecting the type would have worked for pooled features and raised Cannot cast ... to an EmbOptimType for sequence features. Registering a placeholder FTRL class in torchrec's class-to-type table serves both paths from one mechanism. The class never steps -- the dynamicemb kernel replaces the in-backward optimizer with EmptyFusedOptimizer -- so it mirrors torchrec's own placeholders in torchrec/optim/optimizers.py.

Optimizer-state multiplier. FTRL keeps linear and accum per element, i.e. 2 * embedding_dim per row like Adam. torchrec's _get_optimizer_multipler maps any class it does not know to 1, which would have undersized local_hbm_for_values and the planner's HBM/DDR estimate by one embedding width per row, so dynamicemb_util wraps it.

Rejecting non-dynamicemb tables. ftrl_optimizer cannot drive a fused TBE at all. Left alone, the failure is an FBGEMM assertion during DistributedModelParallel construction that names neither dynamicemb nor the offending table. The plan now raises with the table name instead. data_parallel tables are exempt -- they are replicated and updated by dense_optimizer, so the sparse optimizer never reaches them.

initial_accumulator_value needed no new plumbing -- create_sparse_optimizer already pops it into the module-level switch, and TZRecDynamicEmbParameterSharding.get_additional_fused_params already re-injects it per table.

Test Plan

All run in a conda env matching requirements/runtime.txt (torch 2.13.0+cu130) with the new dynamicemb wheel installed.

• tzrec/optim/optimizer_builder_test.py::test_create_sparse_optimizer_ftrl -- the config yields the FTRL class, torchrec resolves it to DynamicEmbOptimType.FTRL, every knob survives into the kwargs, initial_accumulator_value lands in the module-level switch, and the class constructs with those kwargs the way apply_optimizer_in_backward does.
• tzrec/utils/dynamicemb_util_test.py::OptimizerMultiplerTest -- FTRL is 2, and SGD / Adam / RowWiseAdagrad / untrained still delegate to torchrec unchanged.
• tzrec/utils/plan_util_test.py::test_sharding_plan_sizes_ftrl_optimizer_state -- pins the four new kwargs on BatchedDynamicEmbeddingTablesV2.__init__, then checks the planned local_hbm_for_values matches the multiplier-2 sizing (it differs from the multiplier-1 value, so the assertion is load-bearing).
• tzrec/utils/plan_util_test.py::test_sharding_plan_rejects_ftrl_on_non_dynamicemb_table -- a FUSED table raises, a DENSE (data_parallel) one does not.
• tzrec/tests/rank_integration_test.py::test_multi_tower_din_with_dynamicemb_ftrl_train_eval -- new GPU-scoped train + eval on a DIN config whose id features are all dynamicemb-backed, covering both the EmbeddingBagCollection and EmbeddingCollection paths (the latter being where the kwargs approach would have broken). The tables deliberately differ in max_capacity, score_strategy (TIMESTAMP / STEP / LFU / NO_EVICTION) and whether they set init_capacity_per_rank, and one carries a frequency_admission_strategy, so FTRL runs against several distinct table setups, all four score-assignment strategies and a deterministically-firing admission path rather than one uniform configuration. (Capacity-driven eviction is not reachable from a static config in this harness: mock ids are drawn from [0, max_capacity), so the distinct-key count can never exceed capacity.)
• Separately confirmed end to end that the config reaches FTRLDynamicEmbeddingOptimizer inside the table with {'opt_type': 'ftrl', 'lr': 0.01, 'learning_rate_power': -0.5, 'ftrl_beta': 1.0, 'initial_accumulator_value': 0.1, 'l1_reg': 0.01, 'l2_reg': 0.02} and a per-row state dim of 2 * embedding_dim.
• python tzrec/tests/run.py --scope gpu (261 tests), plus the dynamicemb-adjacent unit suites (checkpoint_util, delta_embedding_dump, export_util, modules.embedding, plan_util, dynamicemb_util, optim).
• After merging master (which brought [feat] support probabilistic admission strategy for dynamicemb #680's probabilistic admission strategy and the same wheel pin): the merged dynamicemb_util_test and both dynamicemb integration tests pass.
• pre-commit run -a and pyrefly check clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Qk8SFFzLveowmKz7fePXL1

[github-actions]

gradient_clipping / max_gradient look like dead knobs for this message. ftrl_optimizer can only drive dynamicemb tables (enforced at plan time in dynamicemb_util._to_sharding_plan), but dynamicemb's FTRL path never reads either key: FTRLDynamicEmbeddingOptimizer.get_opt_args/set_opt_args handle exactly {opt_type, lr, learning_rate_power, ftrl_beta, initial_accumulator_value, l1_reg, l2_reg}, and neither the update wrappers nor the CUDA kernels (ftrl_update_for_padded_buffer / ftrl_update_for_flat_table) take clipping args — the base class's clip_gradient/need_gradient_clipping helpers are never called by any optimizer.

Unlike the sibling Fused*Optimizer messages, where these fields reach FBGEMM's OptimizerArgs for plain tables, here BatchedDynamicEmbeddingTablesV2.__init__ silently absorbs them, so gradient_clipping: true is a silent no-op on every table this optimizer can ever drive.

Since the message is new in this PR, consider dropping both fields, or keeping them with a comment marking them unsupported and warning in create_sparse_optimizer when they're set to non-defaults.

[github-actions]

The cross-reference here doesn't land: feature.md does not document which feature types support dynamicemb — its only mention of dynamicemb in the whole file is the negative note that CombineFeature does not support it (line 247). The authoritative list lives in feature.proto (IdFeature / ComboFeature / LookupFeature / MatchFeature / RegexReplaceFeature / CustomFeature / BoolMaskFeature carry the dynamicemb field; RawFeature / ExprFeature / OverlapFeature / TokenizeFeature / CombineFeature do not).

A reader deciding whether their features can use ftrl_optimizer will follow this link and find nothing — consider listing the supported feature types in this note (or adding a positive dynamicemb list to feature.md).

[github-actions]

Nit: this setter docstring now says the switch is used "by Adagrad and, on dynamicemb tables, by FTRL", but the paired getter sparse_init_accumulator_value() (line 141) still reads "Sparse Adagrad accumulator initial value". Worth syncing the one-liner.

[github-actions]

The eviction coverage this config aims for (per the PR description) is coincidental rather than by design. For dynamicemb features the mock generator draws ids from [0, max_capacity) (IdMockInput(num_ids=feature.num_embeddings), and num_embeddings = max_capacity for dynamicemb — id_feature.py:76-77), and test_train_eval produces batch_size x num_parts x 4 = 131,072 rows over TEST_NPROC_PER_NODE=2 ranks:

• id_1 (TIMESTAMP, 200k keyspace): ~56k-96k distinct keys vs 100k capacity per rank — eviction can never fire.
• id_2 (STEP, 65,536), id_3 (LFU, 32,768), id_4_emb (NO_EVICTION, 20,480): the 50-id sequence column saturates the keyspace, so distinct keys ~= max_capacity and per-shard occupancy lands exactly at capacity — whether eviction (or NO_EVICTION insert rejection) actually happens comes down to hash-distribution luck and the RNG seed.

delta_embedding_dump_test.py already has the deterministic pattern: deliberately shrink max_capacity (e.g. 1024-4096) on one evicting table so capacity < distinct keys on every run. Consider sizing at least one table that way so the eviction paths are exercised reliably. (Frequency admission on id_2 is fine — ~50 occurrences per key clears threshold: 2 deterministically.)

[github-actions]

Minor: num_buckets is dead config on dynamicemb features — the table keyspace comes from max_capacity (id_feature.py:76-77) and mock ids are drawn from [0, max_capacity). The existing multi_tower_din_fg_dynamicemb_mock.config omits num_buckets on its dynamicemb feature; keeping num_buckets: 1000000 next to max_capacity: 200000 here is misleading (same for the other four dynamicemb features).

[github-actions]

Minor test hygiene: this call runs register_ftrl_emb_opt_type(), which permanently inserts FTRL into torchrec's global _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE. The class tearDown resets the accumulator switch but not this map, and tzrec/tests/run.py executes the whole suite in one process — so for every later in-process test FTRL appears "registered", which would mask a missing register_ftrl_emb_opt_type() call in a future test. Consider _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE.pop(FTRL, None) in tearDown.

[github-actions]

The exemption is keyed on the compute kernel, but its rationale (the comment here and the new optimizer.md note) is scoped to data_parallel — and DENSE != data_parallel. plan_util._filter_compute_kernels (lines 1208-1212) only drops DENSE when FUSED is also allowed, so non-DP DENSE-kernel options survive on CPU topologies and whenever the user pins compute_kernels: ["dense"] via embedding constraints. An FTRL-configured table planned that way passes with no error or warning, and since the DENSE kernel's params carry no in-backward optimizer, in_backward_optimizer_filter (main.py:931) hands them to the dense optimizer — the table silently trains with a different update rule and LR schedule, which is exactly the class of quiet failure this check was added to prevent, and it contradicts the new doc's promise that any non-dynamicemb sparse table errors at plan time.

Consider keying the exemption on sharding_option.sharding_type == ShardingType.DATA_PARALLEL.value instead (DP always implies the DENSE kernel, so the intended exemption is unchanged) — then CPU / constrained dense-kernel FTRL plans hit the intended, actionable ValueError.

[github-actions]

Code review summary

Static review across five areas (code quality, performance, test coverage, documentation accuracy, security), cross-checked against the installed torchrec 1.8.0 sources and upstream dynamicemb (NVIDIA/recsys-examples PR #487, matching the pinned wheel 0.1.0+20260920.9643985).

The core design checks out. Verified independently:

• The placeholder-class + _OPTIMIZER_CLASS_TO_EMB_OPT_TYPE registration is the right mechanism: torchrec 1.8.0 reads that dict at call time in optimizer_type_to_emb_opt_type (unconditionally on the EmbeddingCollection path), the in-place mutation is visible to all call sites, and the class mirrors torchrec's own placeholders exactly.
• The version guard works as documented: DynamicEmbOptimType was introduced by the upstream FTRL PR itself, so any pre-FTRL wheel (e.g. the previously pinned +20260911) fails at the import and hits the friendly RuntimeError — no AttributeError gap in practice.
• Multiplier 2.0 is correct for FTRL's linear+accum state, both sizing call sites go through the new wrapper, no stale shard_estimators._get_optimizer_multipler call remains in the repo, and torchrec's internal estimator only runs for non-dynamicemb tables (where FTRL is rejected), so there is no leftover undersizing path. Without the fix, HBM/host sizing would have been ~33% short — this was a real OOM risk.
• The plan-time ValueError cannot corrupt the planner search: torchrec calls to_sharding_plan once on the final best plan, outside the PlannerError proposal loop.
• delta_embedding_dump slices values[:, :emb_dim], so FTRL's wider rows won't leak optimizer state into dumps.
• The wheel bump is a pure version-string change on the same first-party host, consistent across all 3 CUDA variants x 3 Python versions and the docs; no leftover references to the old build.

Findings posted inline (2 medium, rest low/minor):

1. optimizer.proto — gradient_clipping/max_gradient are dead knobs on FusedFTRLOptimizer: upstream's FTRL optimizer and CUDA kernels never read them, and the table's **kwargs swallows them silently.
2. dynamicemb_util.py — the plan-time rejection exempts any DENSE-kernel table, but DENSE != data_parallel (CPU topologies, user-pinned compute_kernels: ["dense"]): FTRL there silently trains with the dense optimizer, contradicting the new doc note.
3. optimizer.md — the "可配置dynamicemb的特征类型见特征文档" link doesn't land: feature.md doesn't document which feature types support dynamicemb.
4. ftrl mock config — eviction coverage is coincidental: mock ids are drawn from [0, max_capacity) and the capacities put every evicting table exactly at the saturation boundary, so whether eviction fires depends on hash luck; delta_embedding_dump_test's shrunk-capacity pattern would make it deterministic.
5. ftrl mock config — num_buckets is dead/misleading on dynamicemb features (keyspace is max_capacity); the existing fg dynamicemb config omits it.
6. optimizer.py — getter docstring still says "Adagrad" after the setter's was updated for FTRL.
7. optimizer_builder_test.py — the new test permanently registers FTRL in torchrec's global map without tearDown cleanup (full suite runs in one process).

Nothing blocking; items 1, 2 and 4 are the ones worth addressing before merge.

---

Automated review (5 specialized reviewers + orchestrator verification). All five review areas completed.