python -m perceiver.scripts.text.clm_fsdp fit \
  --model.num_self_attention_layers=20 \
  --model.max_latents=512 \
  --model.num_channels=1280 \
  --model.num_heads=10 \
  --model.max_heads_parallel=2 \
  --model.cross_attention_dropout=0.0 \
  --model.post_attention_dropout=0.0 \
  --model.output_norm=true \
  --model.output_bias=false \
  --model.abs_pos_emb=false \
  --model.init_scale=0.02 \
  --model.optimizer=AdamW \
  --model.optimizer.lr=3e-4 \
  --model.scheduler=CosineWithWarmupLR \
  --model.scheduler.warmup_steps=1000 \
  --model.scheduler.min_fraction=0.1 \
  --model.max_grad_norm=1.0 \
  --data=C4DataModule \
  --data.tokenizer=xlnet-base-cased \
  --data.padding_side=left \
  --data.max_seq_len=1024 \
  --data.min_seq_len=512 \
  --data.batch_size=256 \
  --data.num_train_workers=2 \
  --data.num_valid_workers=1 \
  --trainer.strategy=fsdp_perceiver_ar \
  --trainer.accelerator=gpu \
  --trainer.devices=8 \
  --trainer.precision=bf16 \
  --trainer.max_steps=50000 \
  --trainer.accumulate_grad_batches=1 \
  --trainer.check_val_every_n_epoch=null \
  --trainer.val_check_interval=500 \
  --trainer.limit_val_batches=20 \
  --trainer.log_every_n_steps=20 \
  --trainer.logger=TensorBoardLogger \
  --trainer.logger.save_dir=logs \
  --trainer.logger.name=clm-fsdp
