test

Author: ruoyu61

pretrain/shakespeare.py

@@ -42,14 +42,14 @@
 grad_clip = 1.0
 
 # For shakespeare, choose smaller block size than vanilla LLaMA
-block_size = 1024
+block_size = 128
 
 
 def main() -> None:
     auto_wrap_policy = partial(transformer_auto_wrap_policy, transformer_layer_cls={Block})
     strategy = FSDPStrategy(auto_wrap_policy=auto_wrap_policy, activation_checkpointing=Block, limit_all_gathers=True)
 
-    fabric = L.Fabric(accelerator="cuda", devices=4, precision="bf16-mixed", strategy=strategy)
+    fabric = L.Fabric(accelerator="cuda", devices=1, precision="bf16-true", strategy=strategy)
     fabric.launch()
     fabric.seed_everything(1337 + fabric.global_rank)