setup

EleutherAI · Jun 19, 2024 · 46904d5 · 46904d5
1 parent 4c7cb11
commit 46904d5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/configs/local_setup.yml b/configs/local_setup.yml
@@ -22,6 +22,10 @@
   "load": "checkpoints",
   "checkpoint_validation_with_forward_pass": False,
 
+
+ # "launcher": "openmpi",
+  #"deepspeed_mpi": true,
+
   "tensorboard_dir": "tensorboard",
   "log_dir": "logs",
   "use_wandb": True,

diff --git a/megatron/model/rwkv/v6/rwkv.py b/megatron/model/rwkv/v6/rwkv.py
@@ -223,6 +223,7 @@ def forward(self, x):
         H = self.neox_args.num_attention_heads
 
         r, k, v, g, w = self.jit_func(x)
+        print(f"shape of r: {r.size()}, k: {k.size()}, v: {v.size()}, g: {g.size()}, w: {w.size()}, H: {H}, B: {B}, T: {T}, C: {C}, time_faaaa: {self.time_faaaa.size()},  \n")
         x = RUN_CUDA_RWKV(B, T, C, H, r, k, v, w, u=self.time_faaaa)
 
         return self.jit_func_2(x, g)