[tx] support long sequence

[raulchen]

I'm currently testing / optimizing skyrl-tx for long-sequence use cases. E.g., the tinker search tool recipe

Using this issue to track the related fixes and optimizations.

[raulchen]

TX Memory Benchmark Results

Model: Qwen/Qwen3-4B-Instruct-2507 | TP: 8 | GPUs: 8x A100-40GB

---

Optimization Stack (Cumulative)

Stage	Optimizations
0	baseline (main)
1	+ flash-attention #879
2	+ skip-full-logits #878
3	+ per-layer-checkpointing #906
4	+ chunked-lm-head #902

What Each Optimization Does

1. Flash Attention

• Uses cuDNN flash attention instead of naive attention
• Memory: O(T) instead of O(T²). At seq=4k, avoids materializing 4096×4096 attention matrix per head
• Universal win with no tradeoffs

2. Skip Full Logits

• During prefill, only compute logits for last token instead of all tokens
• Memory: Avoids [batch, seq, vocab] tensor. At BS=16, seq=4k, vocab=152k = 20GB in fp16
• Only helps sampling (prefill). No effect on training

3. Per-Layer Checkpointing

• Use jax.lax.scan with checkpointing so XLA keeps only ONE layer's activations, recomputing during backward
• Memory: 36 layers → 1 layer activation storage (~36x reduction)
• Key training optimization. Trades compute for memory

4. Chunked LM Head

• Compute cross-entropy loss in chunks instead of materializing full logits tensor
• Memory: At BS=16, seq=4k, vocab=152k: 10GB → 300MB (with chunk_size=512)

---

Sample Mode Results

Format: peak memory (MiB) / time (s)

Seq Len = 4096

Stage	BS=16	BS=32	BS=64	BS=128	BS=256
0 (main)	9,345 / 149s	17,537 / 158s	31,397 / 169s	ERROR	-
1 (flash)	9,347 / 149s	17,537 / 152s	31,397 / 165s	ERROR	-
2 (+skip-logits)	5,251 / 150s	9,345 / 149s	21,635 / 158s	31,401 / 172s	ERROR
3 (+checkpoint)	5,251 / 156s	9,345 / 157s	21,635 / 167s	31,401 / 182s	ERROR
4 (all)	5,251 / 155s	9,345 / 158s	21,635 / 166s	31,401 / 177s	ERROR

Seq Len = 8192

Stage	BS=16	BS=32	BS=64	BS=128
0 (main)	19,587 / 177s	ERROR	-	-
1 (flash)	17,539 / 177s	31,397 / 195s	ERROR	-
2 (+skip-logits)	9,347 / 172s	13,441 / 182s	31,399 / 204s	ERROR
3 (+checkpoint)	9,347 / 178s	13,441 / 186s	31,401 / 214s	ERROR
4 (all)	9,347 / 178s	13,441 / 186s	31,399 / 214s	ERROR

---

Train Mode Results

Format: peak memory (MiB) / time (s)

Seq Len = 4096

Stage	BS=1	BS=2	BS=4	BS=8	BS=16	BS=32
0 (main)	31,467 / 180s	ERROR	-	-	-	-
1 (flash)	21,713 / 198s	31,475 / 167s	ERROR	-	-	-
2 (+skip-logits)	21,713 / 198s	31,475 / 167s	ERROR	-	-	-
3 (+checkpoint)	9,313 / 29s	9,313 / 32s	13,409 / 35s	25,697 / 40s	ERROR	-
4 (all)	9,313 / 29s	9,313 / 31s	13,409 / 31s	21,601 / 34s	31,367 / 38s	ERROR

Seq Len = 8192

Stage	BS=1	BS=2	BS=4	BS=8	BS=16
0 (main)	ERROR	-	-	-	-
1 (flash)	31,471 / 180s	ERROR	-	-	-
2 (+skip-logits)	31,471 / 180s	ERROR	-	-	-
3 (+checkpoint)	9,313 / 31s	13,409 / 33s	25,697 / 40s	ERROR	-
4 (all)	9,313 / 31s	13,409 / 32s	21,601 / 34s	31,367 / 39s	ERROR

---

Max Batch Size Summary

Stage	Sample 4k	Sample 8k	Train 4k	Train 8k
0 (main)	64	16	1	0
1 (flash)	64	32	2	1
2 (+skip-logits)	128	64	2	1
3 (+checkpoint)	128	64	8	4
4 (all)	128	64	16	8

---

Summary

With all optimizations enabled:

• Sample: 2x max batch size (64 → 128 at 4k seq, 16 → 64 at 8k seq)
• Train: 16x max batch size (1 → 16 at 4k seq, 0 → 8 at 8k seq)

---

Test Commands

Code #892

# Sample benchmark
uv run --extra tinker python benchmarks/benchmark_memory.py \
    --experiment-name <name>_sample \
    --mode sample \
    --batch-sizes 16,32,64,128,256,512,1024 \
    --seq-lens 4096,8192 \
    --output-root /mnt/user_storage/tx_memory_benchmark

# Train benchmark
uv run --extra tinker python benchmarks/benchmark_memory.py \
    --experiment-name <name>_train \
    --mode train \
    --batch-sizes 1,2,4,8,16,32,64,128,256,512,1024 \
    --seq-lens 4096,8192 \
    --output-root /mnt/user_storage/tx_memory_benchmark

[raulchen]

Beyond the above the above optimizations:

• Train: main memory bottleneck is the layer outputs, can be addressed by sequence parallelism.
• Inference: main memory bottleneck is the KV cache, can be addressed by prefix caching, quantization, etc.