Simplify dtype and device settings (#57)

Author: gty111

gllm/layers/layernorm.py

@@ -9,14 +9,12 @@ def __init__(
         self,
         hidden_size: int,
         eps: float,
-        dtype: torch.dtype = None,
     ) -> None:
         super().__init__()
         self.variance_epsilon = eps
         self.variance_size_override = None
         self.hidden_size = hidden_size
-        self.weight = nn.Parameter(torch.ones(
-            hidden_size, dtype=dtype))
+        self.weight = nn.Parameter(torch.ones(hidden_size))
         self.has_weight = True
 
     def forward(

gllm/layers/rotary_embedding.py

@@ -15,7 +15,6 @@ def __init__(
         max_position_embeddings: int,
         base: float,
         is_neox_style: bool,
-        dtype,
     ) -> None:
         super().__init__()
         self.head_size = head_size
@@ -26,7 +25,7 @@ def __init__(
 
         cache = self._compute_cos_sin_cache()
 
-        cache = cache.to(dtype=dtype,device='cuda')
+        cache = cache.to(dtype=torch.get_default_dtype())
         self.register_buffer("cos_sin_cache", cache, persistent=False)
 
     def _compute_inv_freq(self, base):
@@ -84,14 +83,13 @@ def __init__(
         max_position_embeddings: int,
         base: int,
         is_neox_style: bool,
-        scaling_factors: Union[List[float], float],
-        dtype: torch.dtype,
+        scaling_factors: Union[List[float], float]
     ) -> None:
         if isinstance(scaling_factors, float):
             scaling_factors = [scaling_factors]
         self.scaling_factors: List[float] = scaling_factors  # noqa
         super().__init__(head_size, rotary_dim, max_position_embeddings, base,
-                         is_neox_style, dtype)
+                         is_neox_style)
         # Lazy initialized.
         self._scaling_factor_to_offset: Dict[float, int]
 
@@ -142,7 +140,6 @@ def __init__(
         max_position_embeddings: int,
         base: int,
         is_neox_style: bool,
-        dtype: torch.dtype,
         scaling_factor: float,
         low_freq_factor: float,
         high_freq_factor: float,
@@ -153,7 +150,7 @@ def __init__(
         self.high_freq_factor = high_freq_factor
         self.orig_max_position = orig_max_position
         super().__init__(head_size, rotary_dim, max_position_embeddings, base,
-                         is_neox_style, dtype)
+                         is_neox_style)
 
     def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
         inv_freqs = super()._compute_inv_freq(base)

gllm/model_loader.py

@@ -20,14 +20,6 @@ def __init__(self, load_format, model_path):
         self.model_path = model_path
         self.load_config()
         self.load_format = load_format
-        
-    def get_dtype(self, dtype: str):
-        if dtype == 'float16':
-            return torch.float16
-        elif dtype == 'bfloat16':
-            return torch.bfloat16
-        else:
-            assert 0
 
     def get_finish_tokens(self):
         return self.get_model_type().get_finish_tokens(self.config)

gllm/model_runner.py

@@ -39,7 +39,7 @@ def init(self, mp_load_progress=None):
         memory_manager_cls = PrefixMemoryManager if self.enable_prefix_caching else MemoryManager
         self.memory_manager = memory_manager_cls(
             gpu_memory_util=self.gpu_memory_util, num_layers=self.model.num_layers,
-            dtype=self.model.dtype, page_size=self.page_size, kv_head_num=self.model.num_kv_heads,
+            dtype=self.model_loader.dtype, page_size=self.page_size, kv_head_num=self.model.num_kv_heads,
             kv_head_dim=self.model.head_dim, vocab_size=self.model_loader.vocab_size)
 
     def encode(self, content, chat: bool = False):

gllm/models/chatglm.py

@@ -25,17 +25,17 @@ def __init__(self, layer_id: int, config):
         self.scaling = self.head_dim**-0.5
 
         self.rotary_emb = RotaryEmbedding(
-            self.head_dim, self.head_dim // 2, config.seq_length, getattr(config,'rope_theta',10000), False, config.torch_dtype)
+            self.head_dim, self.head_dim // 2, config.seq_length, getattr(config,'rope_theta',10000), False)
         self.attn = FlashAttention(
             layer_id, self.scaling, self.num_heads, self.num_kv_heads, self.head_dim, self.hidden_size)
 
         self.projection_size = config.kv_channels * self.num_heads
         self.qkv_hidden_size = self.projection_size + 2 * \
             self.head_dim * config.multi_query_group_num
         self.query_key_value = nn.Linear(self.hidden_size, self.qkv_hidden_size,
-                                         bias=config.add_bias_linear or config.add_qkv_bias, dtype=config.torch_dtype, device='cuda')
+                                         bias=config.add_bias_linear or config.add_qkv_bias)
         self.dense = nn.Linear(self.projection_size, self.hidden_size,
-                               bias=config.add_bias_linear, dtype=config.torch_dtype, device='cuda')
+                               bias=config.add_bias_linear)
 
     def forward(self, input_data: InputData, hidden_states: torch.Tensor):
         qkv = self.query_key_value(hidden_states)
@@ -51,10 +51,10 @@ def __init__(self, config):
         super().__init__()
         self.add_bias = config.add_bias_linear
         self.dense_h_to_4h = nn.Linear(
-            config.hidden_size, config.ffn_hidden_size*2, bias=self.add_bias, dtype=config.torch_dtype, device='cuda')
+            config.hidden_size, config.ffn_hidden_size*2, bias=self.add_bias)
         self.activation_func = SiluAndMul()
         self.dense_4h_to_h = nn.Linear(
-            config.ffn_hidden_size, config.hidden_size, bias=self.add_bias, dtype=config.torch_dtype, device='cuda')
+            config.ffn_hidden_size, config.hidden_size, bias=self.add_bias)
 
     def forward(self, hidden_states):
         # [s, b, 4hp]
@@ -73,13 +73,13 @@ def __init__(self, layer_id, config):
 
         assert config.rmsnorm
         self.input_layernorm = RMSNorm(
-            config.hidden_size, config.layernorm_epsilon, config.torch_dtype)
+            config.hidden_size, config.layernorm_epsilon)
 
         self.self_attention = GLMAttention(layer_id, config)
         self.hidden_dropout = config.hidden_dropout
 
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, config.layernorm_epsilon, config.torch_dtype)
+            config.hidden_size, config.layernorm_epsilon)
 
         self.mlp = GLMMLP(config)
 
@@ -127,7 +127,7 @@ def __init__(self, config):
                 assert config.rmsnorm
                 layer_norm_func = RMSNorm
                 self.final_layernorm = layer_norm_func(
-                    config.hidden_size, config.layernorm_epsilon, config.torch_dtype)
+                    config.hidden_size, config.layernorm_epsilon)
 
     def forward(self, input_data: InputData, hidden_states: torch.Tensor):
         for layer in self.layers:
@@ -145,14 +145,14 @@ def __init__(self, config):
         super().__init__()
 
         self.embedding = nn.Embedding(
-            config.padded_vocab_size, config.hidden_size, dtype=config.torch_dtype, device='cuda')
+            config.padded_vocab_size, config.hidden_size)
 
         self.multi_query_group_num = config.multi_query_group_num
         self.kv_channels = config.kv_channels
 
         self.encoder = GLMTransformer(config)
         self.output_layer = nn.Linear(
-            config.hidden_size, config.padded_vocab_size, bias=False, dtype=config.torch_dtype, device='cuda')
+            config.hidden_size, config.padded_vocab_size, bias=False)
 
     def forward(self, input_data: InputData, hidden_states=None):
         if get_pp_rank() == 0:
@@ -169,7 +169,6 @@ def __init__(self, config):
 
         self.config = config
         self.max_model_len = config.seq_length
-        self.dtype = config.torch_dtype
         self.num_kv_heads = config.multi_query_group_num
         self.head_dim = config.hidden_size // config.num_attention_heads
         self.transformer = ChatGLMModel(config)

gllm/models/llama.py

@@ -20,9 +20,9 @@ def __init__(self, config):
         self.hidden_size = config.hidden_size
         self.intermediate_size = config.intermediate_size
         self.gate_up_proj = nn.Linear(
-            self.hidden_size, self.intermediate_size*2, bias=False, dtype=config.torch_dtype, device='cuda')
+            self.hidden_size, self.intermediate_size*2, bias=False)
         self.down_proj = nn.Linear(
-            self.intermediate_size, self.hidden_size, bias=False, dtype=config.torch_dtype, device='cuda')
+            self.intermediate_size, self.hidden_size, bias=False)
         self.act_fn = SiluAndMul()
 
     def forward(self, x: torch.Tensor):
@@ -39,9 +39,9 @@ def __init__(self, layer_id: int, config):
         self.num_key_value_heads = config.num_key_value_heads
 
         self.qkv_proj = nn.Linear(self.hidden_size, (self.num_heads+self.num_key_value_heads*2)
-                                  * self.head_dim, bias=False, dtype=config.torch_dtype, device='cuda')
+                                  * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.num_heads*self.head_dim,
-                                self.hidden_size, bias=False, dtype=config.torch_dtype, device='cuda')
+                                self.hidden_size, bias=False)
 
         self.rope_theta = getattr(config,'rope_theta',10000)
 
@@ -55,18 +55,18 @@ def __init__(self, layer_id: int, config):
                     "original_max_position_embeddings"]
                 self.rotary_emb = Llama3RotaryEmbedding(
                     self.head_dim, self.head_dim, original_max_position,
-                    self.rope_theta, True, config.torch_dtype,
-                    rope_scaling['factor'], low_freq_factor, high_freq_factor, original_max_position)
+                    self.rope_theta, True, rope_scaling['factor'], low_freq_factor, 
+                    high_freq_factor, original_max_position)
             elif rope_scaling['type'] == 'linear':
                 self.rotary_emb = LinearScalingRotaryEmbedding(
                     self.head_dim, self.head_dim, config.max_position_embeddings,
-                    self.rope_theta, True, rope_scaling['factor'], config.torch_dtype)
+                    self.rope_theta, True, rope_scaling['factor'])
             else:
                 assert 0
         else:
             self.rotary_emb = RotaryEmbedding(
                 self.head_dim, self.head_dim, config.max_position_embeddings,
-                self.rope_theta, True, config.torch_dtype)
+                self.rope_theta, True)
 
         self.scaling = self.head_dim**-0.5
 
@@ -90,10 +90,10 @@ class LlamaDecoderLayer(nn.Module):
     def __init__(self, layer_id: int, config):
         super().__init__()
         self.input_layernorm = RMSNorm(
-            config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+            config.hidden_size, config.rms_norm_eps)
         self.self_attn = LlamaAttention(layer_id, config)
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+            config.hidden_size, config.rms_norm_eps)
         self.mlp = LlamaMLP(config)
 
     def forward(self, input_data: InputData, hidden_states: torch.Tensor, residual: Optional[torch.Tensor]):
@@ -128,10 +128,10 @@ def __init__(self, config):
             layer_id-self.start_layer, config) for layer_id in range(self.start_layer, self.end_layer)])
         if get_pp_rank() == 0:
             self.embed_tokens = nn.Embedding(
-                config.vocab_size, config.hidden_size, dtype=config.torch_dtype, device='cuda')
+                config.vocab_size, config.hidden_size)
         if get_pp_rank() == get_pp_size() - 1:
             self.norm = RMSNorm(
-                config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+                config.hidden_size, config.rms_norm_eps)
 
     def forward(self, input_data: InputData, hidden_states=None, residual=None):
         if get_pp_rank() == 0:
@@ -151,7 +151,6 @@ class LlamaForCausalLM(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.max_model_len = config.max_position_embeddings
-        self.dtype = config.torch_dtype
         self.num_kv_heads = config.num_key_value_heads
         self.head_dim = config.hidden_size // config.num_attention_heads
         self.model = LlamaModel(config)
@@ -160,7 +159,7 @@ def __init__(self, config):
         self.ret_residual = True
         if get_pp_rank() == get_pp_size() - 1:
             self.lm_head = nn.Linear(
-                config.hidden_size, config.vocab_size, bias=False, dtype=config.torch_dtype, device='cuda')
+                config.hidden_size, config.vocab_size, bias=False)
         self.sampler = Sampler()
 
     def forward(self, input_data: InputData, hidden_states=None, residual=None):

gllm/models/qwen2.py

@@ -48,7 +48,7 @@ def __init__(self, layer_id: int, config):
             self.hidden_size, (self.num_heads+self.num_kv_heads*2)*self.head_dim, bias=True)
         self.o_proj = nn.Linear(self.num_heads*self.head_dim, self.hidden_size, bias=False)
         self.rotary_emb = RotaryEmbedding(
-            self.head_dim, self.head_dim, self.max_position_embeddings, self.rope_theta, True, config.torch_dtype)
+            self.head_dim, self.head_dim, self.max_position_embeddings, self.rope_theta, True)
         self.attn = FlashAttention(
             layer_id, self.scaling, self.num_heads, self.num_kv_heads, self.head_dim, self.hidden_size)
 
@@ -67,9 +67,9 @@ def __init__(self, layer_id: int, config, attention_type=Qwen2Attention, mlp_typ
         self.self_attn = attention_type(layer_id, config)
         self.mlp = mlp_type(config)
         self.input_layernorm = RMSNorm(
-            config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+            config.hidden_size, config.rms_norm_eps)
         self.post_attention_layernorm = RMSNorm(
-            config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+            config.hidden_size, config.rms_norm_eps)
 
     def forward(self, input_data: InputData, hidden_states: torch.Tensor, residual: Optional[torch.Tensor]):
         if residual is None:
@@ -102,7 +102,7 @@ def __init__(self, config, decoder_layer_type=Qwen2DecoderLayer):
         ])
         if get_pp_rank() == get_pp_size() - 1:
             self.norm = RMSNorm(
-                config.hidden_size, config.rms_norm_eps, config.torch_dtype)
+                config.hidden_size, config.rms_norm_eps)
 
     def forward(self, input_data: InputData, hidden_states=None, residual=None):
         if get_pp_rank() == 0:

gllm/models/qwen3.py

@@ -29,11 +29,11 @@ def __init__(self, layer_id, config):
             self.hidden_size, (self.num_heads+self.num_kv_heads*2)*self.head_dim, bias=self.qkv_bias)
         self.o_proj = nn.Linear(self.num_heads*self.head_dim, self.hidden_size, bias=False)
         self.rotary_emb = RotaryEmbedding(
-            self.head_dim, self.head_dim, config.max_position_embeddings, self.rope_theta, True, config.torch_dtype)
+            self.head_dim, self.head_dim, config.max_position_embeddings, self.rope_theta, True)
         self.attn = FlashAttention(
             layer_id, self.scaling, self.num_heads, self.num_kv_heads, self.head_dim, self.hidden_size)
-        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps, config.torch_dtype)
-        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps, config.torch_dtype)
+        self.q_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, config.rms_norm_eps)
 
     def forward(self, input_data: InputData, hidden_states: torch.Tensor):
         qkv = self.qkv_proj(hidden_states)