Commit 14384d8
committed
Update on "Remove in-place mutations from UMA spin embeddings"
Target configuration: uma-s-1p2 with `torch.compile(dynamic=True)`,
`external_graph_gen=True`, `merge_mole=True`, and `execution_mode="umas_fast_gpu"`.
The checkpoint used by the current benchmark selects `rand_emb`, so the changed
sinusoidal and linear embedding branches are not executed in that measured run.
This fix applies to UMA configurations selecting `pos_emb` or `lin_emb`.
The positional path used one-argument `torch.where`, equivalent to `nonzero`,
whose result length depends on the number of zero-spin entries, followed by
an in-place indexed write. The linear path similarly modified its
input tensor in place, changing its version counter and visible aliases.
```
zero_idxs = torch.where(x == 0)[0]
emb[zero_idxs] = 0
emb = emb * (x != 0).unsqueeze(-1)
```
Likewise, `x[x == 0] = -100` becomes
`x = torch.where(x == 0, -100, x)`. These functional forms preserve null-spin
semantics without data-dependent index shapes or input/intermediate mutation.
**Activation**
No independent flag enables this fix. It is used automatically by UMA configurations whose checkpoint selects the sinusoidal (`pos_emb`) or linear (`lin_emb`) spin embedding. Compilation is enabled with:
```python
settings = InferenceSettings(compile=True)
```
The `uma-s-1p2` checkpoint used by the current performance benchmark selects `rand_emb`, so it does not execute the changed branches. No inference setting switches an existing checkpoint between these embedding implementations.
Test Plan:
```
PYTHONPATH=$PWD/src:$PYTHONPATH pytest -q tests/core/models/uma/nn/test_embedding.py -k test_spin_embedding_fullgraph_preserves_input
```
Authored with assistance from Codex.
[ghstack-poisoned]0 file changed
0 commit comments