misc

lanlan777 · Mar 19, 2023 · 107f167 · 107f167
1 parent 6d4dec7
commit 107f167
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/RWKV-v4neo/src/dataset.py b/RWKV-v4neo/src/dataset.py
@@ -167,6 +167,8 @@ def worker_seed():
                                 magic_prime = 162165671
                             elif args.ctx_len == 4096:
                                 magic_prime = 81082817
+                            elif args.ctx_len == 8192:
+                                magic_prime = 40541399
                             data = self.data_pile
                         else:
                             ii = ii // 2

diff --git a/RWKV-v4neo/train.py b/RWKV-v4neo/train.py
@@ -166,13 +166,11 @@
         elif args.ctx_len == 4096:
             args.magic_prime = 81082817
             args.epoch_count = 2010
+        elif args.ctx_len == 8192:
+            args.magic_prime = 40541399
+            args.epoch_count = 1005
         if args.my_pile_shift < 0:
-            if args.ctx_len == 1024:
-                args.my_pile_shift = 0
-            elif args.ctx_len == 2048:
-                args.my_pile_shift = 512
-            elif args.ctx_len == 4096:
-                args.my_pile_shift = 768
+            args.my_pile_shift = 0
 
         if magic_prime_bak > 0:
             args.magic_prime = magic_prime_bak