From d4b9396084835046da701165bf2a577ab3097ba3 Mon Sep 17 00:00:00 2001
From: "yuze.zyz" <yuze.zyz@alibaba-inc.com>
Date: Thu, 29 Jun 2023 16:08:42 +0800
Subject: [PATCH] Revert "support device_map"

This reverts commit 4d0367718d1b7087f9d4167872807fa066d69fdf.
---
 examples/pytorch/chatglm6b/finetune.py                     | 7 ++-----
 modelscope/models/nlp/chatglm/text_generation.py           | 2 --
 modelscope/models/nlp/chatglm2/text_generation.py          | 2 --
 .../trainers/hooks/checkpoint/checkpoint_processor.py      | 3 ++-
 modelscope/trainers/trainer.py                             | 2 +-
 5 files changed, 5 insertions(+), 11 deletions(-)

diff --git a/examples/pytorch/chatglm6b/finetune.py b/examples/pytorch/chatglm6b/finetune.py
index b39c4a26..2dc85f2a 100644
--- a/examples/pytorch/chatglm6b/finetune.py
+++ b/examples/pytorch/chatglm6b/finetune.py
@@ -192,7 +192,7 @@ if config['model']['type'] == 'chatglm6b':
     model_config['model']['prefix_projection'] = args.prefix_projection
 
 tokenizer = ChatGLMTokenizer.from_pretrained(model_dir, trust_remote_code=True)
-model = Model.from_pretrained(model_dir, cfg_dict=model_config, device_map='auto')
+model = Model.from_pretrained(model_dir, cfg_dict=model_config)
 
 if args.ptuning_checkpoint is not None:
     # Evaluation
@@ -213,7 +213,7 @@ if args.pre_seq_len is not None:
     # P-tuning v2
     model = model.half()
     model.transformer.prefix_encoder.float()
-elif not args.use_lora:
+else:
     # Finetune
     model = model.float()
 
@@ -367,8 +367,6 @@ model.gradient_checkpointing_enable()
 if config['model']['type'] == 'chatglm6b':
     model.enable_input_require_grads()
 
-# import torch
-# model = torch.nn.DataParallel(model).cuda()
 trainer = Seq2SeqTrainer(
     model=model,
     cfg_file='./configuration.json',
@@ -377,7 +375,6 @@ trainer = Seq2SeqTrainer(
     seed=args.seed,
     data_collator=data_collator,
     remove_unused_data=True,
-    device='cpu',
     cfg_modify_fn=cfg_modify_fn)
 trainer.tokenizer = tokenizer
 trainer.train()
diff --git a/modelscope/models/nlp/chatglm/text_generation.py b/modelscope/models/nlp/chatglm/text_generation.py
index 8c663cfa..ff32c86d 100644
--- a/modelscope/models/nlp/chatglm/text_generation.py
+++ b/modelscope/models/nlp/chatglm/text_generation.py
@@ -1267,8 +1267,6 @@ class ChatGLMForConditionalGeneration(ChatGLMPreTrainedModel):
             shift_labels = labels[..., 1:].contiguous()
             # Flatten the tokens
             loss_fct = CrossEntropyLoss(ignore_index=-100)
-            if shift_logits.device != shift_labels.device:
-                shift_labels = shift_labels.to(shift_logits.device())
             loss = loss_fct(
                 shift_logits.view(-1, shift_logits.size(-1)),
                 shift_labels.view(-1))
diff --git a/modelscope/models/nlp/chatglm2/text_generation.py b/modelscope/models/nlp/chatglm2/text_generation.py
index 2f5c3771..be744f14 100644
--- a/modelscope/models/nlp/chatglm2/text_generation.py
+++ b/modelscope/models/nlp/chatglm2/text_generation.py
@@ -1003,8 +1003,6 @@ class ChatGLM2ForConditionalGeneration(ChatGLMPreTrainedModel):
             shift_labels = labels[..., 1:].contiguous()
             # Flatten the tokens
             loss_fct = CrossEntropyLoss(ignore_index=-100)
-            if shift_logits.device != shift_labels.device:
-                shift_labels = shift_labels.to(shift_logits.device())
             loss = loss_fct(
                 shift_logits.view(-1, shift_logits.size(-1)),
                 shift_labels.view(-1))
diff --git a/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py b/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py
index 6a59b0d8..f28fc397 100644
--- a/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py
+++ b/modelscope/trainers/hooks/checkpoint/checkpoint_processor.py
@@ -26,6 +26,7 @@ class CheckpointProcessor:
             trainer: The trainer instance.
             output_dir: The target folder used in inference.
         """
+        model = trainer.unwrap_module(trainer.model)
         config = trainer.cfg
 
         # override pipeline by tasks name after finetune done,
@@ -38,7 +39,7 @@ class CheckpointProcessor:
             config['pipeline'] = {'type': config['task']}
 
         self.copy_files_and_dump_config(trainer, output_dir, config,
-                                        '*.bin')
+                                        self._bin_file(model))
 
     @staticmethod
     def copy_files_and_dump_config(trainer, output_dir, config, bin_file):
diff --git a/modelscope/trainers/trainer.py b/modelscope/trainers/trainer.py
index 2e62be89..d8d87826 100644
--- a/modelscope/trainers/trainer.py
+++ b/modelscope/trainers/trainer.py
@@ -232,7 +232,7 @@ class EpochBasedTrainer(BaseTrainer):
             # A logic to fit the current code
             # Put a DDPHook in if launcher is provided.
             if 'hooks' not in self.cfg.train:
-                self.cfg.train['hooks'] = []
+                self.cfg.train['hooks'] = ConfigDict([])
             self.cfg.train['hooks'].append({
                 'type': 'DDPHook',
                 'launcher': self.launcher