Describe the bug
When tuning the chunk size, the tuner can sometimes re-test a chunk size it already determined to be not viable. This is reproducible with a configuration where max chunk size is 1024 and the max viable chunk size is 128. The non-viable chunk size 256 gets tested twice.
To Reproduce
I'll send a PR shortly that includes a unit test for this.
Expected behavior
Chunk size tuning should add minimal overhead and avoid re-testing already tested sizes.
Stack trace
I detected this because I added some verbose logging:
Tuning chunk size for TemplatePairBlock between 4 and 1024
Chunk size 1024 is viable for TemplatePairBlock
Selected chunk size 1024 for TemplatePairBlock
Tuning chunk size for MSAModuleBlock between 4 and 1024
Chunk size 1024 is not viable for MSAModuleBlock
Chunk size 64 is viable for MSAModuleBlock
Chunk size 256 is not viable for MSAModuleBlock
Chunk size 128 is viable for MSAModuleBlock
Chunk size 256 is not viable for MSAModuleBlock
Selected chunk size 128 for MSAModuleBlock
Tuning chunk size for PairFormerBlock between 4 and 1024
Chunk size 1024 is not viable for PairFormerBlock
Chunk size 64 is viable for PairFormerBlock
Chunk size 256 is viable for PairFormerBlock
Chunk size 512 is not viable for PairFormerBlock
Selected chunk size 256 for PairFormerBlock
Tuning chunk size for DiffusionConditioning._forward between 4 and 2048
Chunk size 2048 is viable for DiffusionConditioning._forward
Selected chunk size 2048 for DiffusionConditioning._forward
Describe the bug
When tuning the chunk size, the tuner can sometimes re-test a chunk size it already determined to be not viable. This is reproducible with a configuration where max chunk size is 1024 and the max viable chunk size is 128. The non-viable chunk size 256 gets tested twice.
To Reproduce
I'll send a PR shortly that includes a unit test for this.
Expected behavior
Chunk size tuning should add minimal overhead and avoid re-testing already tested sizes.
Stack trace
I detected this because I added some verbose logging: