Add FTRL-Proximal as a fused TBE optimizer (CUDA + CPU)

[tiankongdeguiji]

Adds OptimType.FTRL (FTRL-Proximal) as a fused TBE optimizer, with both CUDA and CPU backends.

FTRL-Proximal is the standard optimizer for the sparse/wide part of CTR models, and is what Alibaba's x-deeplearning uses for its embedding parameter server. There was no equivalent in FBGEMM, so recsys workloads that rely on it have nothing to migrate onto.

Algorithm

Per element, with state accum (running sum of squared gradients) and linear:

new_accum = accum + g^2
sigma_new = new_accum ^ (-learning_rate_power)
sigma_old = accum     ^ (-learning_rate_power)
linear   += g - (sigma_new - sigma_old) / learning_rate * weight
quadratic = sigma_new / learning_rate + 2 * l2_reg
weight    = |linear| > l1_reg ? (l1_reg * sgn(linear) - linear) / quadratic : 0
accum     = new_accum

This is bit-for-bit the same formulation as TensorFlow's ApplyFtrlV2 (with l2_shrinkage = 0) and x-deeplearning's FtrlUpdater, so checkpoints and hyperparameters carry over from either. learning_rate_power == -0.5 (the default) is special-cased to sqrtf/std::sqrt, which is both faster and what x-deeplearning does.

Two properties worth calling out, both inherent to the algorithm rather than artifacts of this implementation:

• FTRL recomputes the weight from (accum, linear) instead of incrementing it, so a row's initial random value is discarded the first time that row is touched.
• The update divides by learning_rate, so learning_rate > 0 is required. The TBE constructor asserts this for FTRL rather than letting it produce inf/NaN weights.

State layout

momentum1 = accum, momentum2 = linear, both elementwise (D floats per row), the same footprint as ADAM. get_optimizer_state() exposes them under TensorFlow's slot names, accum and linear.

Knobs

Three new TBE constructor arguments, all defaulted so they are no-ops for every other optimizer:

arg	default	meaning
ftrl_learning_rate_power	-0.5	exponent applied to accum; -0.5 gives the classic 1/sqrt schedule
ftrl_l1_reg	0.0	L1 strength; a coordinate whose |linear| never exceeds this stays exactly 0, which is what makes FTRL sparsify embedding rows
ftrl_l2_reg	0.0	proximal L2 strength

They are ftrl_-prefixed because OptimizerArgs is a flat namespace shared by every optimizer. weight_decay is deliberately not reused for ftrl_l2_reg: FBGEMM's weight_decay is half of a (weight_decay, weight_decay_mode) pair with six modes and enters as a gradient term, whereas FTRL's L2 enters the denominator as 2 * l2 and honors none of those modes. Overloading it would make weight_decay_mode a silent no-op on FTRL tables.

initial_accumulator_value is intentionally not implemented — no FBGEMM optimizer has a state-init argument today (Adagrad included), so an FTRL-only one would be inconsistent. A generic optimizer_state_init_values TBE argument would be the right way to add it later.

Support matrix

CUDA	yes
CPU	yes
VBE	yes
global weight decay	no — FTRL's proximal L2 is not a gradient-space decay, so GWD does not compose with it
SSD / optimizer offloading	no (the EmbOptimType metadata arms are filled in, so enabling it later is mechanical)

Tests

fbgemm_gpu/test/tbe/training/backward_optimizers_test.py:

• test_backward_optimizers_ftrl — the shared randomized harness over variable T/D/B/E/L, mixed dims, weighted/unweighted, VBE and all pooling modes, on CPU and CUDA. The new reference block recomputes accum, linear and the post-step weights from the formula above and compares against split_optimizer_states() / split_embedding_weights().
• FtrlTest — a deterministic class covering what the randomized harness cannot check cheaply: ftrl_l1_reg producing exactly 0.0, the generic powf path agreeing with the reference (and differing from the sqrt fast path), the 2 * l2_reg denominator term, multi-step accum/linear accumulation, zero-gradient idempotence, and CPU/CUDA parity.

🤖 Generated with Claude Code

https://claude.ai/code/session_01FPQnjDECjd8ZEvGF5Hpzpn

[tiankongdeguiji]

The TorchRec side that consumes this is meta-pytorch/torchrec#4710 — it adds a dense torchrec.optim.FTRL counterpart, the optimizer-class ↔ EmbOptimType mapping, and the planner's optimizer-state multiplier. That PR gates its registration on hasattr(EmbOptimType, "FTRL") so it stays importable against fbgemm_gpu wheels that predate this change.