Conversation
Seeding a row's optimizer state was the storage's job: every call site filled the whole state region with one scalar, `initial_accumulator_value`, taken from a copy cached on DynamicEmbTableState. That is wrong twice over. Adam does not have an `initial_accumulator_value`. Neither torch.optim.Adam nor FBGEMM's TBE exposes a way to seed its moments -- both start them at zero -- because the bias correction 1/(1-beta^t) is derived assuming it, and a non-zero first moment steers the early steps by a phantom momentum instead of the gradient (a negative gradient still moves the weight down). Configure Adam together with that Adagrad-family option today and its m and v are both silently set to it. Adam now reports zero and warns that the option is ignored. The cached scalar also went stale: set_opt_args() runs on checkpoint load and rewrites initial_accumulator_value, but DynamicEmbTableState captured it at table construction, so keys inserted after a load were seeded with the value the table was built with. Every reader now asks the optimizer, and the cached field is gone -- the optimizer is a single instance shared by the cache, both HybridStorage tiers and backward, so there is one source of truth to ask. Callers now go through BaseDynamicEmbeddingOptimizer.reset_optimizer_states(), which owns the layout of the region it writes. This also gives an optimizer whose state regions do not all start at the same value somewhere to say so, which a single scalar cannot express. Its `indices` argument keeps the fused value buffer intact: the state region has to be sliced off first, since `values[rows, max_emb_dim:]` is a copy a write would be lost to, while `values[:, max_emb_dim:]` is a view. Storage keeps init_optimizer_state() for callers that want the value itself rather than a reset; the benchmark scales random data by it. It no longer forwards the reset, so the interface lives in one place. Also renames get/set_initial_optim_states to get/set_initial_optimizer_state: it returns one scalar, not a collection, and every other public name in the file spells out "optimizer". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Implements FTRL-Proximal, Algorithm 1 of McMahan et al., "Ad Click
Prediction: a View from the Trenches" (KDD 2013). Per coordinate:
new_accum = accum + grad^2
linear += grad - (new_accum^-p - accum^-p) / lr * weight
weight = 0 if |linear| <= l1_reg
= (sign(linear)*l1_reg - linear)
/ ((ftrl_beta + new_accum^-p) / lr + l2_reg)
accum = new_accum
`learning_rate`, `ftrl_beta`, `l1_reg` and `l2_reg` are the paper's alpha,
beta, lambda1 and lambda2. The paper fixes the learning rate at
alpha / (beta + sqrt(n)); `learning_rate_power` generalizes the exponent the
way TensorFlow's FtrlOptimizer does, and its default of -0.5 recovers the
paper exactly. Since that default is also the overwhelmingly common case, it
takes a sqrt path rather than a general pow, chosen on the host because the
predicate is launch-uniform.
Verified two ways without a GPU. Lifting update_one out of the header and
compiling it standalone reproduces a closed form worked out by hand -- with
l1 = l2 = beta = 0 the update collapses to
w*(1 - n^0.5/n_next^0.5) - lr*g/n_next^0.5, which for w=1, lr=0.1, n=0.1 and
g=1 gives 0.6031424. And a differential test against a literal transcription
of Algorithm 1, over 1.44M updates across three exponents and four
(beta, l1, l2) combinations, agrees to 2.6e-6 absolute with no disagreement
about which side of the l1 threshold a coordinate falls on. That residual is
float against the reference's double; float is deliberate, since every other
optimizer here computes in float and double is 1/2 to 1/64 throughput on GPU.
FBGEMM's EmbOptimType has no FTRL member and a Python enum with members
cannot be extended, so dynamicemb carries its own `DynamicEmbOptimType`. The
two enums never compare equal, which is the behaviour we want: an
`== EmbOptimType.X` test elsewhere must not match FTRL.
Per row the state is `linear` then `accum`, each emb_dim wide, the same shape
of layout Adam uses for m and v. `linear` starts at zero and `accum` at
`initial_accumulator_value`, so reset_optimizer_states gains `emb_dims` to
locate the boundary -- a padded buffer reserves the widest table's state for
every row, so the block can be wider than a given row's own state and the
width alone cannot be divided up. Optimizers whose state is uniform ignore it.
Note that seeding `accum` is not free here the way it is in the linear
regression FTRL was written for, whose weights start at zero. It leaves the
state inconsistent with the weight already in the row, and the first update
reconciles them by shrinking the weight by sqrt(n0 / (n0 + g^2)) -- for the
small gradients typical of embeddings, nearly all of it. `ftrl_beta` bounds
the early steps without that cost, since beta sits outside the accumulator;
DynamicEmb_APIs.md tabulates the difference.
construct_twin_module now takes every width of a value row from the table
rather than mixing in the twin's own `dim`, and says so if they disagree.
FTRL has no FBGEMM counterpart, so a twin cannot be built for it anyway.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…s own methods truncate_optimizer_states_for_checkpoint and pad_optimizer_states_from_checkpoint both took the optimizer as their first argument and asked it for the widths they convert between, so they were methods already. They are now states_for_checkpoint / states_from_checkpoint on BaseDynamicEmbeddingOptimizer -- named for what they hand back rather than for how they get there, since truncating and padding turned out to be one optimizer's business rather than everyone's. That is the other half of the change. The conversion only ever applied to rowwise Adagrad, the one optimizer whose get_ckpt_state_dim differs from its get_state_dim, because it widens its single accumulator to a fixed 16 bytes for alignment in the fused value row. Every other optimizer checkpoints exactly what it runs with. So the narrowing and widening moves into an override on RowWiseAdaGradDynamicEmbeddingOptimizer and the base does the common thing: hand the state back for a checkpoint, and change only its precision on the way in. It also drops a redundant parameter: `pad` asked its caller for the scalar to fill the columns a narrow checkpoint does not cover, and all four callers passed optimizer.get_initial_optimizer_state() -- a value the method can reach itself. Those columns now go through reset_optimizer_states, which seeds them the way a fresh row is seeded instead of assuming one scalar covers the whole region. Behaviour change: every one of the four paths now rejects a block that is not the width it keeps, where the old pair quietly truncated a too-wide one and padded a too-narrow one. Each direction has exactly one legal width -- runtime going out, checkpoint coming in -- so any other width is a checkpoint that does not belong to this table, and reshaping it into something that loads only defers the damage to training. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
FtrlVecOptimizer described its exponent through two indirections, neither of which survived reading. It held -learning_rate_power under the name `neg_lr_power`, so the member did not match the option it comes from and following the kernel meant carrying a sign flip in your head. It now holds `learning_rate_power` as given and negates at the single place the exponent is applied, where the negate folds away because the value is uniform. And the flag selecting the fast path was called `lr_power_is_half`, which names the wrong number: learning_rate_power is -0.5, and what is a half is the exponent applied to the accumulator. It is `use_sqrt` now, after what it chooses -- sqrtf over a general powf. The test for that flag is exact rather than within 1e-6. -0.5 is representable in float and arrives that way from Python, so the tolerance never rescued a value that meant to be a square root; it only stood ready to swallow one that did not, silently computing a different exponent than the caller asked for. TensorFlow compares exactly here for the same reason. Both paths were checked to agree: forcing a learning_rate_power of -0.5 down the powf branch instead of the sqrtf one moves the weight by at most 3e-7 over 800k updates. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
The accumulator is raised to -learning_rate_power, so a positive value moves it into the denominator of the learning rate: the rate then grows as the gradients accumulate instead of decaying, and training walks off. Zero is meaningful -- accum**0 is 1, a fixed learning rate -- so the bound is <= 0, which is the same one Keras enforces on its Ftrl. Checked in the constructor and again in set_opt_args, since that takes a checkpoint's meta and nothing else vets it. The learning-rate check that was already there moves alongside it so both entry points share one rule. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
FTRLVecOptimizer and SGDVecOptimizer, to read the same way as the FTRLDynamicEmbeddingOptimizer and SGDDynamicEmbeddingOptimizer they implement. Adam, AdaGrad and RowWiseAdaGrad keep their spelling -- those are names rather than initialisms. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…ayout test/unit_test.sh names the files it runs one by one, so test_ftrl_optimizer.py was never going to execute. Added to the fwd_bwd list. test_padded_buffer_optimizer.py covers the corruption path that issue NVIDIA#419 found in Adam: the padded buffer reserves the widest table's state for every row, so a narrower table's second state region begins at max_emb_dim + edim rather than at 2 * max_emb_dim. FTRL has exactly the same exposure with linear and accum, and had no test for it. The new case runs two tables of different dims through ftrl_update_for_padded_buffer -- across both the vec4 and the scalar path, as the Adam cases do -- and checks the narrow table's accum advanced by g^2 in its own slot, that linear was computed from that accum rather than from a zero read out of the padding, and that nothing past the row's own state was written. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
test_alignment_cache_storage_shapes checks that a table's value row is embedding_dim + get_optimizer_state_dim wide once sharding and bucketing have had their say. FTRL is the one optimizer whose state is two regions rather than one, so its row is 3 * dim where every other swept optimizer's is 1x or 2x -- exactly the kind of width this test exists to pin down, and it was not being swept. Nothing else on the path needed widening: the planner does not read the optimizer, and the sharder passes fused_params through untouched, so DynamicEmbOptimType reaches the table factory as-is. The two optimizer_type annotations become OptimType to admit it. This grows the case count by a third (1080 -> 1440). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
FTRL had a file to itself because nothing in the suite could serve as its reference: the batched-table tests check optimizer arithmetic against a SplitTableBatchedEmbeddingBagsCodegen built with the same hyperparameters, and FBGEMM has no FTRL. A second dynamicemb storage is no substitute -- PyDictStorage takes the very optimizer under test, so a wrong formula comes out wrong on both sides and the comparison passes. So give the suite the reference it was missing. optimizer_reference.py transcribes Algorithm 1 in float64, outside any test_ module so pytest leaves it alone, and test_ftrl_backward_matches_reference drives a real BatchedDynamicEmbeddingTablesV2 against it over four iterations -- sequence pooling with one unique key per bag and loss.sum(), so every row is touched once and each gradient is exactly 1. It sweeps beta, L1, L2, the exponent and a seeded accumulator, across a vec4 and a non-vec4 dim, and checks all three of weight, linear and accum. That subsumes the closed-form anchor, the float64 comparison and the exact-zero check that lived in the standalone file. The two cases that were never module-level behaviour move to test_padded_buffer_optimizer.py, which is already where the kernels' own contracts are pinned: that the flat-table and padded-buffer kernels agree, and that a row whose index is -1 is left alone. Dropped: the state-width case, now covered twice over by test_alignment and test_dump_optimizer_states_ckpt_width, along with the reset-semantics and hyperparameter-validation cases. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
test/unit_tests/optimizer/ now holds both: test_padded_buffer_optimizer.py moves in unchanged, and test_ftrl_optimizer.py joins it carrying the three tests that are about FTRL rather than about a buffer layout -- the two fused kernels agreeing with each other, the -1 row-addressing contract, and a real BatchedDynamicEmbeddingTablesV2 landing where the paper says it should. ftrl_step comes along as the reference those last tests need, so the separate optimizer_reference.py module is gone; a reference with one caller belongs next to it. test_ftrl_padded_buffer_mixed_dims_accum stays behind, since what it pins is the padded buffer's layout rather than FTRL's arithmetic, and it reads alongside the two Adam cases that cover the same ground. Also restores a @cuda marker that the previous move dropped from test_ftrl_flat_table_matches_padded_buffer, which would have made it error rather than skip on a machine without a GPU. Both files are named in unit_test.sh at their new paths. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
test_ftrl_backward_matches_reference walked export_keys_values_iter and checked each batch against the full reference. With sixteen keys and a default batch size of 65536 there is only ever one batch, so it passed -- but by coincidence, not by construction: a larger table or a smaller batch size would have compared a subset against the whole and failed on shape. The per-batch argsort had the same flaw, ordering only within a batch when the reference is indexed by key across all of them. The keys are 0..num_keys-1 and the reference is indexed the same way, so a batch can just look itself up. That drops the sort and the concatenation, and stops the test caring what order the export walks the table in or how it divides the rows. A running count still has to reach num_keys, so a short export cannot pass either. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
jiashuy
changed the title
Collaborator
Author
|
/build |
Contributor
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Collaborator
|
❌ Pipeline #68097371 -- failed
Result: 9/17 jobs passed |
_create_optimizer builds OptimizerArgs from four names it never receives -- learning_rate_power, ftrl_beta, l1_reg and l2_reg are parameters of __init__, and the method is a separate scope. That construction happens before the optimizer type is dispatched on, so every configuration hit it, not just FTRL: building a BatchedDynamicEmbeddingTablesV2 with SGD, Adam or Adagrad raised NameError too. Declare the four on _create_optimizer and forward them from the call site. Also drops three imports left unused when the code that wanted them went away. Both this and the NameError are what pyflakes reports and py_compile does not; the suite now runs clean against main on the files this branch touches. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…izer Three errors, all in the direction of making a non-zero seed look safer than it is. The exponent convention contradicted itself: the formula block defines p = learning_rate_power and raises the accumulator to n^(-p), so the paper's square root is p = -0.5, not the p = 0.5 the prose claimed -- a value the implementation rejects outright. The retention table was wrong by two orders of magnitude at the small end. For n0 = 0.1 and g = 0.0001 the initializer keeps 0.000005%, not the 0.0005% printed. And the recommendation was backwards. ftrl_beta was offered as the safe way to bound the early steps, on the grounds that it sits outside the accumulator. It does, but it discounts the initializer just the same: with n0 = 0 the retained fraction is |g| / (ftrl_beta + |g|), which for beta = 1 and g = 0.01 is under one percent. Both knobs are now given by one formula, (sqrt(n1) - sqrt(n0)) / (ftrl_beta + sqrt(n1)), the table covers their four combinations, and the text says plainly that only leaving both at zero keeps the initializer whole -- and what that costs in turn, a first step of w - lr*sign(g) whatever the gradient's size. The optimizer docstring says the same. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68127510 -- failed
Result: 3/5 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68295263 -- failed
Result: 14/17 jobs passed |
test_ftrl_backward_matches_reference splats one opt_params dict into two places, and they do not take the same keys. The table constructor accepts initial_accumulator_value; ftrl_step does not, because seeding the accumulator is not part of a step -- the reference already starts from the seeded value, having read ref_accum out of the table after reset_optimizer_states laid it down. The seeded_accum case was the only one carrying that key, so it was the only one that raised TypeError. Split the step's own arguments out before the call. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68348321 -- failed
Result: 3/5 jobs passed |
Collaborator
Author
|
/build |
Collaborator
|
❌ Pipeline #68373272 -- failed
Result: 14/17 jobs passed |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Checklist