update README

OpenBMB · Apr 25, 2024 · 227d9c6 · 227d9c6
1 parent 8ef00c5
commit 227d9c6
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -70,7 +70,6 @@ Before you can use BMTrain, you need to initialize it at the beginning of your c
 import bmtrain as bmt
 bmt.init_distributed(
     seed=0,
-    zero_level=3,   # support 2 and 3 now
     # ...
 )
 ```
@@ -120,9 +119,9 @@ class MyModule(bmt.DistributedModule): # changed here
         super().__init__()
         self.param = bmt.DistributedParameter(torch.empty(1024)) # changed here
         self.module_list = torch.nn.ModuleList([
-            bmt.Block(SomeTransformerBlock()), # changed here
-            bmt.Block(SomeTransformerBlock()), # changed here
-            bmt.Block(SomeTransformerBlock())  # changed here
+            bmt.Block(SomeTransformerBlock(), zero_level=3), # changed here, support 2 and 3 now
+            bmt.Block(SomeTransformerBlock(), zero_level=3), # changed here, support 2 and 3 now
+            bmt.Block(SomeTransformerBlock(), zero_level=3)  # changed here, support 2 and 3 now
         ])
 
     def forward(self):