Disable flash encoder latent rotary#789
Conversation
fomo-bot
force-pushed
the
agent/disable-flash-encoder-latent-rotary
branch
from
496cc3d to
791a7aa
Compare
| assert cfg.data.loading.num_workers == 2 | ||
|
|
||
|
|
||
| def test_flash_encoder_perceiver_does_not_enable_latent_rotary(monkeypatch): |
There was a problem hiding this comment.
@fomo-bot I don't think we need a test for this.
Can you also change the the naive perceiver to have equivalent behavior here if it does not already?
There was a problem hiding this comment.
Pushed e9405bf8 to the PR branch.
Changes:
- Removed the added config test per review.
- Updated the naive encoder Perceiver path to use the same explicit
FourierFeatures2Dwrapper as flash. - Set naive
fourier_encode_data=Falseand adjustedinput_channelsto avoid double-encoding. - Updated the Fourier feature docstring to reflect shared behavior.
Verification:
uv run --with pre-commit pre-commit run --files src/samudra/config.py src/samudra/models/modules/augment_input.py tests/test_config.pyuv run pytest tests/test_config.py tests/test_encoder.py- Naive
PerceiverConfig.build(..., implementation="naive")forward smoke check passed.
Debug info
- Repository:
m2lines/Samudra - Trigger: Disable flash encoder latent rotary #789 (comment)
- Request ID:
39d135f8-fe05-454e-9c9d-b06ef336528b - Accepted:
true
There was a problem hiding this comment.
@fomo-bot why are you still passing num_freq_bands & max_freq then?
There was a problem hiding this comment.
Because perceiver_pytorch.Perceiver.__init__ requires num_freq_bands and max_freq as keyword-only args even when fourier_encode_data=False. They are not used to add internal Fourier channels in that mode. I added a short comment in 9252185 so that is explicit at the call site.
There was a problem hiding this comment.
Pushed 9252185d to the PR branch.
I added a short comment explaining that perceiver_pytorch.Perceiver still requires num_freq_bands and max_freq as constructor args even when fourier_encode_data=False.
Also replied on the review thread with that explanation.
Verification passed:
uv run --with pre-commit pre-commit run --files src/samudra/config.pyuv run pytest tests/test_config.py tests/test_encoder.py
Debug info
- Repository:
m2lines/Samudra - Trigger: Disable flash encoder latent rotary #789 (comment)
- Request ID:
39fe6d12-82e1-48db-b1c7-3413c3beaa47 - Accepted:
true
2 participants
Disable the FlashPerceiver encoder's latent rotary embedding by no longer passing
latent_rotary_emb_dimfrom Samudra's config builder. This does not seem likely to be super important, since we still add 2D Fourier features to input patch tokens, and the encoder/decoder still add the physical patch position and scale encodings.We want to do this because on ORCD H200 runs,
samudra-multiis crashing (CUDA memory fault) in FlashAttention's Triton rotary kernel. I tried some variants of the configuration which seem to work but I would rather just disable and avoid future related trouble given it seems unnecessary.