For the long-context language modeling task, we set the learning rate to 0.01 by performing small search over values of 0.003, 0.005, 0.01 and 0.03. We used the Adafactor optimizer (Shazeer & Stern, 2018) with linear warmup with 1000 steps, followed by cosine decay. We applied gradient checkpointing after each segment to save memory. The batch size was set to 64. For the LLM experiments, we set the learning rate to 0.0001 during continual pre-training and task fine-tuning.