cleanup

EleutherAI · Nov 6, 2024 · c6fac96 · c6fac96
1 parent 48de682
commit c6fac96
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/configs/rwkv/430M.yml b/configs/rwkv/430M.yml
@@ -12,7 +12,7 @@
   "output_layer_parallelism": "column",
   "norm": "rmsnorm",
   "rms_norm_epsilon": 1.0e-5,
-  "train_micro_batch_size_per_gpu": 4,
+  "train_micro_batch_size_per_gpu": 1,
 
   "attention_config": [[["rwkv"], 24]],
 

diff --git a/megatron/model/rwkv/v6/rwkv.py b/megatron/model/rwkv/v6/rwkv.py
@@ -258,7 +258,6 @@ def forward(self, x):
         B, T, C = x.size()
         C_tp = C//mpu.get_model_parallel_world_size()
         H = self.neox_args.num_attention_heads//mpu.get_model_parallel_world_size()
-        H_tp = H//mpu.get_model_parallel_world_size()
 
         r, k, v, g, w = self.jit_func(x)