Add Support for GPT-2 Training on different Devices (#551)

* fix bfloat16 * update tokenizer encode * support python 3.10- * support python 3.10- * support python 3.10- * click * fix tokenizer path * update * update * checkout main files
Oneflow-Inc · Sep 13, 2024 · 13056f4 · 13056f4
1 parent 7aeb4ca
commit 13056f4
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/libai/models/gpt_model.py b/libai/models/gpt_model.py
@@ -244,7 +244,9 @@ def forward(self, input_ids, past_length=0):
         bsz, seq_length = input_ids.size()
 
         position_ids = self.position_ids[:, past_length : past_length + seq_length]
-        position_ids = position_ids.expand_as(input_ids).to_global(sbp=input_ids.sbp)
+        position_ids = position_ids.expand_as(input_ids).to_global(
+            sbp=input_ids.sbp, placement=input_ids.placement
+        )
 
         token_embeds = self.token_embeddings(input_ids)
         position_embeds = self.position_embeddings(position_ids)

diff --git a/tools/train_net.py b/tools/train_net.py
@@ -36,7 +36,8 @@ def main(args):
 
     seed_for_rank = cfg.train.seed + flow.env.get_rank()
     flow.manual_seed(seed_for_rank)
-    flow.cuda.manual_seed(seed_for_rank)
+    if flow.cuda.is_available():
+        flow.cuda.manual_seed(seed_for_rank)
     np.random.seed(seed_for_rank)
     random.seed(seed_for_rank)