change nn to models (#3032)

applications/ChatGPT/README.md

@@ -41,7 +41,8 @@ Simplest usage:
 ```python
 from chatgpt.trainer import PPOTrainer
 from chatgpt.trainer.strategies import ColossalAIStrategy
-from chatgpt.nn import GPTActor, GPTCritic, RewardModel
+from chatgpt.models.gpt import GPTActor, GPTCritic
+from chatgpt.models.base import RewardModel
 from copy import deepcopy
 from colossalai.nn.optimizer import HybridAdam
 

applications/ChatGPT/benchmarks/benchmark_gpt_dummy.py

@@ -4,7 +4,8 @@ from copy import deepcopy
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from chatgpt.nn import GPTActor, GPTCritic, RewardModel
+from chatgpt.models.base import RewardModel
+from chatgpt.models.gpt import GPTActor, GPTCritic
 from chatgpt.trainer import PPOTrainer
 from chatgpt.trainer.callbacks import PerformanceEvaluator
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy

applications/ChatGPT/benchmarks/benchmark_opt_lora_dummy.py

@@ -4,7 +4,8 @@ from copy import deepcopy
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from chatgpt.nn import OPTActor, OPTCritic, RewardModel
+from chatgpt.models.base import RewardModel
+from chatgpt.models.opt import OPTActor, OPTCritic
 from chatgpt.trainer import PPOTrainer
 from chatgpt.trainer.callbacks import PerformanceEvaluator
 from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy

applications/ChatGPT/chatgpt/experience_maker/base.py

@@ -4,7 +4,7 @@ from typing import Optional
 
 import torch
 import torch.nn as nn
-from chatgpt.nn.actor import Actor
+from chatgpt.models.base import Actor
 
 
 @dataclass

applications/ChatGPT/chatgpt/experience_maker/naive.py

 import torch
-from chatgpt.nn.utils import compute_reward, normalize
+from chatgpt.models.utils import compute_reward, normalize
 
 from .base import Experience, ExperienceMaker
 

applications/ChatGPT/chatgpt/models/__init__.py

+from .base import Actor, Critic, RewardModel
+from .loss import PairWiseLoss, PolicyLoss, PPOPtxActorLoss, ValueLoss
+
+__all__ = ['Actor', 'Critic', 'RewardModel', 'PolicyLoss', 'ValueLoss', 'PPOPtxActorLoss', 'PairWiseLoss']

applications/ChatGPT/chatgpt/models/base/__init__.py

+from .actor import Actor
+from .critic import Critic
+from .reward_model import RewardModel
+
+__all__ = ['Actor', 'Critic', 'RewardModel']

applications/ChatGPT/chatgpt/nn/actor.py → applications/ChatGPT/chatgpt/models/base/actor.py

@@ -4,9 +4,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .generation import generate
-from .lora import LoRAModule
-from .utils import log_probs_from_logits
+from ..generation import generate
+from ..lora import LoRAModule
+from ..utils import log_probs_from_logits
 
 
 class Actor(LoRAModule):

applications/ChatGPT/chatgpt/nn/critic.py → applications/ChatGPT/chatgpt/models/base/critic.py

@@ -3,8 +3,8 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from .lora import LoRAModule
-from .utils import masked_mean
+from ..lora import LoRAModule
+from ..utils import masked_mean
 
 
 class Critic(LoRAModule):

applications/ChatGPT/chatgpt/nn/reward_model.py → applications/ChatGPT/chatgpt/models/base/reward_model.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch
 import torch.nn as nn
 
-from .lora import LoRAModule
+from ..lora import LoRAModule
 
 
 class RewardModel(LoRAModule):

applications/ChatGPT/chatgpt/models/bloom/__init__.py

+from .bloom_actor import BLOOMActor
+from .bloom_critic import BLOOMCritic
+from .bloom_rm import BLOOMRM
+
+__all__ = ['BLOOMActor', 'BLOOMCritic', 'BLOOMRM']

applications/ChatGPT/chatgpt/nn/bloom_actor.py → applications/ChatGPT/chatgpt/models/bloom/bloom_actor.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch
 from transformers import BloomConfig, BloomForCausalLM, BloomModel
 
-from .actor import Actor
+from ..base import Actor
 
 
 class BLOOMActor(Actor):

applications/ChatGPT/chatgpt/nn/bloom_critic.py → applications/ChatGPT/chatgpt/models/bloom/bloom_critic.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 from transformers import BloomConfig, BloomForCausalLM, BloomModel
 
-from .critic import Critic
+from ..base import Critic
 
 
 class BLOOMCritic(Critic):

applications/ChatGPT/chatgpt/nn/bloom_rm.py → applications/ChatGPT/chatgpt/models/bloom/bloom_rm.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch.nn as nn
 from transformers import BloomConfig, BloomForCausalLM, BloomModel
 
-from .reward_model import RewardModel
+from ..base import RewardModel
 
 
 class BLOOMRM(RewardModel):

applications/ChatGPT/chatgpt/models/gpt/__init__.py

+from .gpt_actor import GPTActor
+from .gpt_critic import GPTCritic
+from .gpt_rm import GPTRM
+
+__all__ = ['GPTActor', 'GPTCritic', 'GPTRM']

applications/ChatGPT/chatgpt/nn/gpt_actor.py → applications/ChatGPT/chatgpt/models/gpt/gpt_actor.py

@@ -3,7 +3,7 @@ from typing import Optional
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel
 
-from .actor import Actor
+from ..base import Actor
 
 
 class GPTActor(Actor):

applications/ChatGPT/chatgpt/nn/gpt_critic.py → applications/ChatGPT/chatgpt/models/gpt/gpt_critic.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.gpt2.modeling_gpt2 import GPT2Model
 
-from .critic import Critic
+from ..base import Critic
 
 
 class GPTCritic(Critic):

applications/ChatGPT/chatgpt/nn/gpt_rm.py → applications/ChatGPT/chatgpt/models/gpt/gpt_rm.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from transformers.models.gpt2.configuration_gpt2 import GPT2Config
 from transformers.models.gpt2.modeling_gpt2 import GPT2Model
 
-from .reward_model import RewardModel
+from ..base import RewardModel
 
 
 class GPTRM(RewardModel):