Online weight updates from torch.distributed (#2279)

2024-12-01 23:23:18 -08:00
parent 28bc60dcab
commit 983bfcf386
12 changed files with 1120 additions and 61 deletions
--- a/python/sglang/srt/models/llama.py
+++ b/python/sglang/srt/models/llama.py
@@ -307,6 +307,8 @@ class LlamaForCausalLM(nn.Module):
        self.quant_config = quant_config
        self.torchao_config = global_server_args_dict["torchao_config"]
        self.model = LlamaModel(config, quant_config=quant_config)
+        # Llama 3.2 1B Insturct set tie_word_embeddings to True
+        # Llama 3.1 8B Insturct set tie_word_embeddings to False
        if self.config.tie_word_embeddings:
            self.lm_head = self.model.embed_tokens
        else: