[feat] support initial_accumulator_value on dynamicemb tables

[tiankongdeguiji]

Background

initial_accumulator_value (TensorFlow Adagrad's same-named parameter, which EasyRec users rely on) was already exposed on FusedAdagradOptimizer, but it only reached FBGEMM TBE tables. SplitTableBatchedEmbeddingBagsCodegen has no such kwarg and would raise on the **fused_params splat, so create_sparse_optimizer popped the key out of the optimizer kwargs and instead passed it through an env var consumed by the apply_split_helper patch.

dynamicemb tables read the very same fused params, so the key never reached them and OptimizerArgs.initial_accumulator_value stayed at its 0.0 default. A model with both kinds of tables therefore trained them with different Adagrad semantics, silently.

FusedRowWiseAdagradOptimizer had no such field at all, although both backends support one.

Changes

• Carry the value to dynamicemb tables through get_additional_fused_params() of the customized-kernel parameter sharding. torchrec merges that per table after the optimizer kwargs, dynamicemb strips only its own planner keys, and BatchedDynamicEmbeddingTables takes initial_accumulator_value as a plain kwarg — so no dynamicemb-side change and no new monkeypatch.
• Add initial_accumulator_value to FusedRowWiseAdagradOptimizer. dynamicemb's EXACT_ROWWISE_ADAGRAD honours it, and the existing FBGEMM patch already keys on the momentum1 state prefix, so rowwise adagrad gets it on both backends.
• Move the pop into the shared prologue of create_sparse_optimizer so both Adagrad variants are handled in one place, and wrap the env var in named accessors instead of repeating the FBGEMM string at the new call site.

Alternatives considered

Letting the key stay in the optimizer kwargs and stripping it inside the FBGEMM kernel would have needed a second monkeypatch covering every **fused_params splat site (split TBE, SSD TBE, ...). Sharder-level fused_params were not usable either: DynamicEmbeddingBagCollectionSharder subclasses EmbeddingBagCollectionSharder and shards the whole EBC, so the value would have leaked onto the FBGEMM tables in the same collection. The per-table customized-kernel hook is the only place the two backends are actually distinguishable.

Test Plan

• New optimizer_builder_test case: adagrad / rowwise adagrad / unset — the key never reaches the returned kwargs and the recorded value matches the config.
• New plan_util_test.PlanUtilDynamicEmbE2ETest case: the generated sharding plan's parameter sharding carries the value in its additional fused params, dynamicemb's pop_additional_fused_params leaves it in place, and BatchedDynamicEmbeddingTables still declares the kwarg (so an upstream rename fails loudly instead of silently dropping the value).
• Verified on GPU with a single-rank DistributedModelParallel shard that the real dynamicemb table's OptimizerArgs.initial_accumulator_value is the configured 0.1, for both adagrad_optimizer and rowwise_adagrad_optimizer, and that a plain FBGEMM rowwise-adagrad table's momentum1 buffer is filled with 0.1.
• multi_tower_din_fg_dynamicemb_mock.config now sets initial_accumulator_value: 0.1, so test_multi_tower_din_with_dynamicemb_train_eval exercises the path end to end — passes.
• test_multi_tower_din_with_fg_adagrad_init_acc_train_eval_export (the FBGEMM regression) and all of tzrec/optim, tzrec/utils/plan_util_test, tzrec/utils/dynamicemb_util_test, tzrec/main_test pass.
• pre-commit run on the touched files and pyrefly check are clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_013vLyBqTcb9nAcDbXN37nUs

[github-actions]

The rewritten gate (env var → module global) and the momentum1 fill below it have no direct test coverage: the new optimizer_builder_test case asserts only the kwargs/global plumbing, and the e2e multi_tower_din_fg_mock_adagrad_init_acc.config run asserts success only — a regression in use_init_value or in the torch.full fills would pass CI. Since this PR reworks exactly this mechanism (and extends it to rowwise adagrad), consider a small unit test that invokes apply_split_helper with a fake split / set_attr_fn / persistent_state_fn and asserts the buffer contents for prefix="momentum1".

[github-actions]

This reads the process global at plan time, so correctness depends on create_sparse_optimizer running before any plan generation. _train_and_evaluate satisfies that (main.py:885 before :906), but evaluate / export / predict build plans with the 0.0 default — benign today (state is checkpoint-restored and no updates run), yet the contract is implicit and untested. Consider documenting it in the set_sparse_init_accumulator_value docstring ("must be called before planning / TBE init"), mirroring the set_auto_retain_evicted_keys pattern in this module.

[github-actions]

Two wording issues:

1. “（对齐TensorFlow Adagrad的同名参数，默认0.0）” can be misread as TensorFlow's default being 0.0 — TF's is 0.1 (as the code comment in tzrec/optim/optimizer.py notes). Suggest e.g. “对齐TensorFlow Adagrad的同名参数（TF默认0.1），TZRec默认0.0”.
2. “新插入的key其accumulator会初始化为该值” describes dynamicemb (HKV) semantics only — for regular FBGEMM tables the whole momentum1 buffer is pre-filled at table construction, with no notion of newly inserted keys. Suggest scoping that clause to the dynamicemb表.

[github-actions]

Code review — PR #669

Reviewed with five parallel focus areas (code quality, performance, test coverage, documentation accuracy, multi-process/state safety), plus independent verification of the cross-library plumbing. No blocking issues found — the design is sound.

Verified as correct:

• Ordering: create_sparse_optimizer (main.py:885) always runs before create_planner/collective_plan in _train_and_evaluate, so both readers of the new module global (the FBGEMM apply_split_helper patch at TBE-init time and the dynamicemb plan-time read) see the configured value on every rank.
• Upstream APIs: torchrec add_params_from_parameter_sharding merges get_additional_fused_params() per table for CUSTOMIZED_KERNEL (present in 1.2 through 1.8); the dynamicemb DynamicEmbParameterSharding base is a dataclass over ParameterSharding with all-default fields, and its get_additional_fused_params enumerates fields(DynamicEmbParameterSharding) only — so the subclass override here is genuinely required — while pop_additional_fused_params pops only the three hardcoded planner keys, letting initial_accumulator_value through to BatchedDynamicEmbeddingTablesV2.__init__ and on to OptimizerArgs.
• Checkpoint plan JSON serializes a fixed field whitelist and is only read back for constraints, so the extra dataclass field never round-trips; plans are always regenerated.
• Env-var removal is clean: zero remaining references to FBGEMM_MOMENTUM1_STATE_INIT_VALUE anywhere; init_value != 0.0 is equivalent to the old truthiness gate. config_to_kwargs (including_default_value_fields=True) guarantees the pop/reset runs for both adagrad messages.
• The plan-level test assertions on dynamicemb internals (inspect.signature canary, pop_additional_fused_params behavior) are a sensible guard against silent value drops, since the V2 constructor swallows unknown **kwargs.

Posted 3 inline comments:

1. tzrec/optim/optimizer.py — the reworked momentum1 fill has no direct test asserting buffer contents (e2e tests assert success only).
2. tzrec/utils/dynamicemb_util.py — the "setter must run before planning" contract is implicit; suggest documenting it on the setter.
3. docs/source/models/optimizer.md — the new note ambiguously implies the TF default is 0.0 (it is 0.1), and the newly-inserted-key clause describes dynamicemb semantics but is written to cover both table kinds.

Minor nits (no action required):

• The 3-line inline comment at optimizer_builder.py:51-53 stretches the repo "one short line" comment guidance; the cross-module explanation would fit naturally in the set_sparse_init_accumulator_value docstring instead.
• tzrec/version.py bump: version bumps have often landed as dedicated PRs in this repo — confirm including it here is intentional.

🤖 Generated with Claude Code

[tiankongdeguiji]

Addressed the review in bb959d7.

Fixed

1. optimizer.py — no direct coverage of the momentum1 fill. Real gap, and the one worth closing: this PR both reworked the gate and extended the mechanism to rowwise adagrad. Added SparseInitAccumulatorValueTest in tzrec/optim/optimizer_test.py, covering EXACT_ADAGRAD (accumulator per element) and EXACT_ROWWISE_ADAGRAD (per row) — it asserts the momentum1 buffer contents and width, that weights is left alone (which pins the prefix == "momentum1" gate), and that a 0.0 value leaves the buffer zeroed. Rather than the suggested fake split / callbacks, it builds a real CPU SplitTableBatchedEmbeddingBagsCodegen, so it also proves fbgemm still routes through the patched helper, and it needs no GPU lane. Mutation-checked: forcing the gate false, or flipping the prefix to momentum2, both fail it.

2. dynamicemb_util.py — implicit "set before planning" contract. Documented on set_sparse_init_accumulator_value, mirroring set_auto_retain_evicted_keys. Confirmed the impact note is right and benign: apply_optimizer_in_backward runs only in train_and_evaluate (main.py:889, before create_planner at :894); evaluate and predict_checkpoint build no in-backward optimizers, so their dynamicemb tables are training=False and hold no optimizer state for the default to affect.

3. optimizer.md — both wording issues. The note now reads TF默认0.1，TorchEasyRec默认0.0, and splits the semantics: a regular table fills the whole accumulator at table build time, a dynamicemb table fills a key's accumulator on first write.

4. Nit: the 3-line inline comment. Folded into the setter docstring, leaving one line at the call site. Also dropped the surrounding if, so a non-Adagrad optimizer resets the value instead of inheriting a stale one.

Not changed

5. Nit: the tzrec/version.py bump. Keeping it — the last five commits touching that file (18d185e, 4950156, c3e9df8, 2a117eb, b30998a) are all ordinary feature/bugfix PRs that bump it in place, and master is at 1.4.7 so 1.4.8 is conflict-free.

Re-ran after the change: the new CPU test, optimizer_builder_test, plan_util_test, dynamicemb_util_test, main_test, and both GPU e2e runs (test_multi_tower_din_with_dynamicemb_train_eval, test_multi_tower_din_with_fg_adagrad_init_acc_train_eval_export) — all green; pre-commit and pyrefly check clean.

🤖 Generated with Claude Code

https://claude.ai/code/session_013vLyBqTcb9nAcDbXN37nUs