Skip to content

Commit a9939cb

Browse files
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]
2 parents 457b7b7 + dc59424 commit a9939cb

0 file changed

File tree

    0 commit comments

    Comments
     (0)