Integrate xlstm cleanly. #35377
Integrate xlstm cleanly. #35377
Conversation
kpoeppel
force-pushed
the
integrate_xlstm_clean
branch
4 times, most recently
from
4a4e347 to
fe5759b
Compare
kpoeppel
requested review from
ArthurZucker,
Rocketknight1,
gante and
stevhliu
as code owners
|
How can this PR keep failing inside other models' tests, how can they get into 'main'? |
| xlstm_block_config = xLSTMLargeConfig( | ||
| vocab_size=config.vocab_size, | ||
| embedding_dim=config.embedding_dim, | ||
| num_blocks=config.num_blocks, | ||
| num_heads=config.num_heads, | ||
| use_bias=config.use_bias, | ||
| add_out_norm=config.add_out_norm, | ||
| norm_eps=config.norm_eps, | ||
| norm_reduction_force_float32=config.norm_reduction_force_float32, | ||
| # mlstm_layer | ||
| qk_dim_factor=config.qk_dim_factor, | ||
| v_dim_factor=config.v_dim_factor, | ||
| # mlstm backend | ||
| chunkwise_kernel=config.chunkwise_kernel, | ||
| sequence_kernel=config.sequence_kernel, | ||
| step_kernel=config.step_kernel, | ||
| mode=config.mode, | ||
| chunk_size=config.chunk_size, | ||
| return_last_states=config.return_last_states, | ||
| autocast_kernel_dtype=config.autocast_kernel_dtype, | ||
| eps=config.eps, | ||
| inference_state_dtype=config.inference_state_dtype, | ||
| # feedforward | ||
| ffn_proj_factor=config.ffn_proj_factor, | ||
| ffn_round_up_to_multiple_of=config.ffn_round_up_to_multiple_of, | ||
| # capping | ||
| gate_soft_cap=config.gate_soft_cap, | ||
| output_logit_soft_cap=config.output_logit_soft_cap, | ||
| weight_mode=config.weight_mode, | ||
| ) |
There was a problem hiding this comment.
we should align xLSTMLargeConfig to match the inputs of mLSTMBlock
There was a problem hiding this comment.
and this not have to do this here
There was a problem hiding this comment.
There are still slight deviations of the xLSTMLargeConfig compare to the xLSTMConfig in configuration_xlstm.py. So I think this conversion is necessary actually.
There was a problem hiding this comment.
Some refactoring needed for the camel casing of classes!
There was a problem hiding this comment.
Is the casing ok as it is now?
superbock
force-pushed
the
integrate_xlstm_clean
branch
3 times, most recently
from
4545076 to
69b9d5c
Compare
superbock
force-pushed
the
integrate_xlstm_clean
branch
from
b925fc6 to
3f6aec9
Compare
|
So, finally we got all tests passing, and have integrated a standalone version in the code now, so the external libraries would just be needed for further speed ups. Thanks for the further review! |
I think the documentation should now be similar, thanks for pointing this out! |
|
@Rocketknight1 @Cyrilvallez , sorry for the long-time stall of this, I think it should be ready now for a review, thanks! :) |
|
cc @Cyrilvallez as I'll be off for 2 weeks! |
What does this PR do?
This PR integrates
xLSTMvia thexlstm-library including certain optimizations (potentially use torch.compile and cuda graphs for speed up). This enables using theNX-AI/xLSTM-7bwithout a special fork of transformers.Before submitting
Pull Request section?
to it if that's the case.
documentation guidelines, and
here are tips on formatting docstrings.
Who can review?
Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@ArthurZucker