Repository: Lightning-AI/litgpt

Issue:
Overwrite system prompt on load/generate
Currently the prompt style is created on model generation and cannot easily be changed.
However it would make sense to update the system prompt for each generation. to offer more flexibility.
While the `PromptStyle` format is defined for each supported model architecture, there is no need to enforce the system prompt text itself.
Therefore I suggest an extension to the `PromptStyle` class to update the system prompt

Discussion:
@rasbt: Thanks for the suggestion. This is a reasonable extension. Unfortunately, I currently don't have the capacity to work on it, but maybe my colleagues can weigh in here.

Solution patch:
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
new file mode 100644
index 0000000..ddb9eea
--- /dev/null
+++ b/.devcontainer/devcontainer.json
@@ -0,0 +1,105 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the README at:
+// https://github.com/microsoft/vscode-dev-containers/tree/v0.194.0/containers/python-3
+{
+  "name": "Python 3 (litgpt)",
+  "build": {
+    "dockerfile": "Dockerfile",
+    "context": "..",
+    "args": {
+      "VARIANT": "3.11-bookworm"
+    }
+  },
+  "runArgs": [
+    // Enable GPU passthrough, requires WSL2 on Windows
+    //"--gpus=all",
+    // One of the following options is required for torch multiprocessing
+    //"--ipc=host",
+    //"--shm-size=4gb",
+  ],
+  // Features to add to the dev container. More info: https://containers.dev/features.
+  "features": {
+    "ghcr.io/devcontainers/features/git:1": {},
+    "ghcr.io/devcontainers/features/git-lfs:1": {},
+    //"ghcr.io/devcontainers/features/nvidia-cuda:1": {},
+    "ghcr.io/devcontainers-extra/features/actionlint:1": {},
+    "ghcr.io/devcontainers-extra/features/pre-commit:2": {},
+    "ghcr.io/dhoeric/features/act:1": {},
+    "ghcr.io/devcontainers/features/docker-in-docker:2": {
+      "version": "latest",
+      "moby": true
+    }
+  },
+  // Set *default* container specific settings.json values on container create.
+  "customizations": {
+    "vscode": {
+      "settings": {
+        "editor.tabSize": 4,
+        "editor.renderWhitespace": "all",
+        "editor.formatOnSave": true,
+        "editor.rulers": [120],
+        "files.exclude": {
+          "**/__pycache__": true
+        },
+        "python.pythonPath": "/usr/local/bin/python",
+        "python.defaultInterpreterPath": "/usr/local/bin/python",
+        "python.languageServer": "Pylance",
+        "python.analysis.autoImportCompletions": true,
+        "python.analysis.completeFunctionParens": true,
+        "python.analysis.autoSearchPaths": true,
+        "python.testing.pytestArgs": ["tests"],
+        "python.testing.unittestEnabled": false,
+        "python.testing.pytestEnabled": true,
+        "code-eol.highlightNonDefault": true,
+        "code-eol.highlightExtraWhitespace": true,
+        "autoDocstring.docstringFormat": "google-notypes",
+        "autoDocstring.guessTypes": true,
+        "autoDocstring.generateDocstringOnEnter": true,
+        "autoDocstring.startOnNewLine": true,
+        "telemetry.telemetryLevel": "off",
+        "[python]": {
+          "editor.formatOnSave": true,
+          "editor.defaultFormatter": "charliermarsh.ruff",
+          "editor.codeActionsOnSave": {
+            "source.organizeImports": "always",
+            "source.fixAll": "always"
+          }
+        }
+      },
+      // Add the IDs of extensions you want installed when the container is created.
+      "extensions": [
+        "ms-python.python",
+        "ms-python.vscode-pylance",
+        "ms-toolsai.jupyter",
+        "GitHub.copilot",
+        "GitHub.copilot-chat",
+        "github.vscode-github-actions",
+        "SanjulaGanepola.github-local-actions",
+        "charliermarsh.ruff",
+        "esbenp.prettier-vscode",
+        "ms-vscode.test-adapter-converter",
+        "njqdev.vscode-python-typehint",
+        "KevinRose.vsc-python-indent",
+        "medo64.render-crlf",
+        "shardulm94.trailing-spaces",
+        "nhoizey.gremlins",
+        "wayou.vscode-todo-highlight",
+        "Gruntfuggly.todo-tree",
+        "njpwerner.autodocstring",
+        "rodolphebarbanneau.python-docstring-highlighter",
+        "mechatroner.rainbow-csv",
+        "uctakeoff.vscode-counter",
+        "bierner.github-markdown-preview",
+        "yahyabatulu.vscode-markdown-alert",
+        "ms-vscode-remote.vscode-remote-extensionpack",
+        "ms-azuretools.vscode-docker",
+        "redhat.vscode-yaml"
+      ]
+    }
+  },
+  // Use 'forwardPorts' to make a list of ports inside the container available locally.
+  // "forwardPorts": [],
+  // Use 'postCreateCommand' to run commands after the container is created.
+  "postCreateCommand": "pre-commit install && pip install '.[extra,compiler,test]' -U",
+  // Comment out connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
+  "remoteUser": "vscode"
+}
diff --git a/README.md b/README.md
index ee640d2..929d78c 100644
--- a/README.md
+++ b/README.md
@@ -53,7 +53,7 @@ Every LLM is implemented from scratch with **no abstractions** and **full contro
 # Quick start
 Install LitGPT
 ```
-pip install 'litgpt[all]'
+pip install 'litgpt[extra]'
 ```
 
 Load and use any of the [20+ LLMs](#choose-from-20-llms):
diff --git a/extensions/xla/generate/adapter.py b/extensions/xla/generate/adapter.py
index 81e5a99..e835834 100644
--- a/extensions/xla/generate/adapter.py
+++ b/extensions/xla/generate/adapter.py
@@ -26,6 +26,7 @@ def setup(
     prompt: str = "What food do llamas eat?",
     *,
     input: str = "",
+    sys_prompt: Optional[str] = None,
     adapter_path: Path = Path("out/adapter/alpaca/lit_model_adapter_finetuned.pth"),
     checkpoint_dir: Path = Path("checkpoints/tiiuae/falcon-7b"),
     max_new_tokens: int = 100,
@@ -40,6 +41,7 @@ def setup(
     Args:
         prompt: The prompt/instruction (Alpaca style).
         input: Optional input (Alpaca style).
+        sys_prompt: Optional system prompt.
         adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
             `xla/finetune/adapter.py`.
         checkpoint_dir: The path to the checkpoint folder with pretrained model weights.
@@ -52,13 +54,14 @@ def setup(
     devices = XLAAccelerator.auto_device_count()
     strategy = XLAFSDPStrategy(auto_wrap_policy={Block}) if devices > 1 else "auto"
     fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy)
-    fabric.launch(main, prompt, input, adapter_path, checkpoint_dir, max_new_tokens, top_k, temperature)
+    fabric.launch(main, prompt, input, sys_prompt, adapter_path, checkpoint_dir, max_new_tokens, top_k, temperature)
 
 
 def main(
     fabric: L.Fabric,
     prompt: str,
     input: str,
+    sys_prompt: Optional[str],
     adapter_path: Path,
     checkpoint_dir: Path,
     max_new_tokens: int,
@@ -90,7 +93,7 @@ def main(
     tokenizer = Tokenizer(checkpoint_dir)
     # TODO: Load prompt style from checkpoint and apply it here
     prompt_style = Alpaca()
-    prompt = prompt_style.apply(prompt, input=input)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt, input=input)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/api.py b/litgpt/api.py
index d2db974..32cc196 100644
--- a/litgpt/api.py
+++ b/litgpt/api.py
@@ -461,6 +461,7 @@ class LLM(torch.nn.Module):
     def generate(
         self,
         prompt: str,
+        sys_prompt: Optional[str] = None,
         max_new_tokens: int = 50,
         temperature: float = 1.0,
         top_k: Optional[int] = None,
@@ -474,6 +475,8 @@ class LLM(torch.nn.Module):
         Arguments:
             model: The model to use.
             prompt: The prompt string to use for generating the samples.
+            sys_prompt: The system prompt string to use for generating the samples.
+                The system prompt allows the user to provide additional instructions to shape all responses by providing additional context, behavioral guidelines, style, and constraints.
             max_new_tokens: The maximum number of new tokens to return.
             temperature: Scales the predicted logits by 1 / temperature.
             top_k: If specified, only sample among the tokens with the k highest probabilities.
@@ -501,7 +504,7 @@ class LLM(torch.nn.Module):
                 "The model is not initialized yet; use the .distribute() "
                 "or .trainer_setup() method to initialize the model."
             )
-        input_ids = self._text_to_token_ids(prompt)
+        input_ids = self._text_to_token_ids(prompt, sys_prompt)
         prompt_length = input_ids.size(0)
         max_returned_tokens = prompt_length + max_new_tokens
 
@@ -564,9 +567,9 @@ class LLM(torch.nn.Module):
         else:
             return self.preprocessor.decode(outputs)
 
-    def _text_to_token_ids(self, prompt):
+    def _text_to_token_ids(self, prompt: str, sys_prompt: Optional[str] = None) -> torch.Tensor:
         """Utility method to convert a prompt text to token IDs"""
-        prompt = self.prompt_style.apply(prompt)
+        prompt = self.prompt_style.apply(prompt, sys_prompt=sys_prompt)
         input_ids = self.preprocessor.encode(prompt)
         return input_ids
 
diff --git a/litgpt/config.py b/litgpt/config.py
index acbea69..70bd607 100644
--- a/litgpt/config.py
+++ b/litgpt/config.py
@@ -9,7 +9,18 @@ import torch
 import yaml
 from typing_extensions import Self
 
-from litgpt.utils import find_multiple
+
+def find_multiple(n: int, k: int) -> int:
+    """Utility function for finding the nearest value to n which is a multiple of k.
+
+    NOTE: We define this function in this module rather than `litgpt.utils` so that users can import
+    this file to do configuration manipulations in Python environments which do not include all the dependencies
+    demanded by `litgpt.utils`.
+    """
+    assert k > 0
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
 
 
 @dataclass
diff --git a/litgpt/deploy/serve.py b/litgpt/deploy/serve.py
index 13ec69b..44aa0b5 100644
--- a/litgpt/deploy/serve.py
+++ b/litgpt/deploy/serve.py
@@ -160,7 +160,8 @@ class OpenAISpecLitAPI(BaseLitAPI):
             config = json.load(fp)
             chat_template = config.get("chat_template", None)
             if chat_template is None:
-                raise ValueError("chat_template not found in tokenizer config file.")
+                print("The tokenizer config does not contain chat_template, falling back to a default.")
+                chat_template = "{% for m in messages %}{{ m.role }}: {{ m.content }}\n{% endfor %}Assistant: "
             self.chat_template = chat_template
 
         self.template = Template(self.chat_template)
diff --git a/litgpt/generate/adapter.py b/litgpt/generate/adapter.py
index d697827..fb7f75c 100644
--- a/litgpt/generate/adapter.py
+++ b/litgpt/generate/adapter.py
@@ -29,6 +29,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     input: str = "",
+    sys_prompt: Optional[str] = None,
     adapter_path: Path = Path("out/finetune/adapter/final/lit_model.pth.adapter"),
     quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"]] = None,
     max_new_tokens: int = 100,
@@ -46,6 +47,7 @@ def main(
         checkpoint_dir: The path to the checkpoint folder with pretrained model weights.
         prompt: The prompt/instruction (Alpaca style).
         input: Optional input (Alpaca style).
+        sys_prompt: Optional system prompt.
         adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
             ``litgpt.finetune.adapter``.
         quantize: Whether to quantize the model and using which method:
@@ -103,7 +105,7 @@ def main(
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
 
-    prompt = prompt_style.apply(prompt, input=input)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt, input=input)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/adapter_v2.py b/litgpt/generate/adapter_v2.py
index 3eb3468..e7a65fa 100644
--- a/litgpt/generate/adapter_v2.py
+++ b/litgpt/generate/adapter_v2.py
@@ -29,6 +29,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     input: str = "",
+    sys_prompt: Optional[str] = None,
     adapter_path: Path = Path("out/finetune/adapter-v2/final/lit_model.pth.adapter_v2"),
     quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"]] = None,
     max_new_tokens: int = 100,
@@ -46,6 +47,7 @@ def main(
         checkpoint_dir: The path to the checkpoint folder with pretrained model weights.
         prompt: The prompt/instruction (Alpaca style).
         input: Optional input (Alpaca style).
+        sys_prompt: Optional system prompt.
         adapter_path: Path to the checkpoint with trained adapter weights, which are the output of
             ``litgpt.finetune.adapter_v2``.
         quantize: Whether to quantize the model and using which method:
@@ -103,7 +105,7 @@ def main(
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
 
-    prompt = prompt_style.apply(prompt, input=input)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt, input=input)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/base.py b/litgpt/generate/base.py
index 552916a..565ef08 100644
--- a/litgpt/generate/base.py
+++ b/litgpt/generate/base.py
@@ -432,6 +432,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     *,
+    sys_prompt: Optional[str] = None,
     num_samples: int = 1,
     max_new_tokens: int = 50,
     top_k: Optional[int] = 50,
@@ -448,6 +449,7 @@ def main(
     Args:
         checkpoint_dir: The checkpoint directory to load.
         prompt: The prompt string to use for generating the samples.
+        sys_prompt: The system prompt to use for generating the samples.
         num_samples: The number of text samples to generate.
         max_new_tokens: The number of generation steps to take.
         top_k: The number of top most probable tokens to consider in the sampling process.
@@ -504,7 +506,7 @@ def main(
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
 
-    prompt = prompt_style.apply(prompt)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/full.py b/litgpt/generate/full.py
index 93ee95c..78cc8dd 100644
--- a/litgpt/generate/full.py
+++ b/litgpt/generate/full.py
@@ -28,6 +28,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     input: str = "",
+    sys_prompt: Optional[str] = None,
     finetuned_path: Path = Path("out/full/alpaca/lit_model_finetuned.pth"),
     quantize: Optional[Literal["bnb.nf4", "bnb.nf4-dq", "bnb.fp4", "bnb.fp4-dq", "bnb.int8"]] = None,
     max_new_tokens: int = 100,
@@ -45,6 +46,7 @@ def main(
         checkpoint_dir: The path to the checkpoint folder with pretrained model weights.
         prompt: The prompt/instruction (Alpaca style).
         input: Optional input (Alpaca style).
+        sys_prompt: Optional system prompt.
         finetuned_path: Path to the checkpoint with trained weights, which are the output of
             ``litgpt.finetune.full``.
         quantize: Whether to quantize the model and using which method:
@@ -101,7 +103,7 @@ def main(
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
 
-    prompt = prompt_style.apply(prompt, input=input)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt, input=input)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/sequentially.py b/litgpt/generate/sequentially.py
index 2ed1987..1aecccb 100644
--- a/litgpt/generate/sequentially.py
+++ b/litgpt/generate/sequentially.py
@@ -138,6 +138,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     *,
+    sys_prompt: Optional[str] = None,
     num_samples: int = 1,
     max_new_tokens: int = 50,
     top_k: Optional[int] = 50,
@@ -154,6 +155,7 @@ def main(
     Args:
         checkpoint_dir: The checkpoint directory to load.
         prompt: The prompt string to use for generating the samples.
+        sys_prompt: The system prompt to use for generating the samples.
         num_samples: The number of text samples to generate.
         max_new_tokens: The number of generation steps to take.
         top_k: The number of top most probable tokens to consider in the sampling process.
@@ -213,7 +215,7 @@ def main(
     prompt_style = (
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
-    prompt = prompt_style.apply(prompt)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/speculative_decoding.py b/litgpt/generate/speculative_decoding.py
index 814ee95..653583b 100644
--- a/litgpt/generate/speculative_decoding.py
+++ b/litgpt/generate/speculative_decoding.py
@@ -331,6 +331,7 @@ def main(
     target_model_checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     *,
+    sys_prompt: Optional[str] = None,
     num_samples: int = 1,
     max_new_tokens: int = 50,
     speculative_k: int = 3,
@@ -349,6 +350,7 @@ def main(
         draft_model: Smaller/faster model used for initial token predictions
         target_model: Larger/more accurate model used to verify draft predictions
         prompt: The prompt string to use for generating the samples.
+        sys_prompt: The system prompt to use for generating the samples.
         num_samples: The number of text samples to generate.
         max_new_tokens: The number of generation steps to take.
         speculative_k: Number of tokens to speculatively generate at each step
@@ -413,7 +415,7 @@ def main(
         if has_prompt_style(target_model_checkpoint_dir)
         else PromptStyle.from_config(target_config)
     )
-    prompt = prompt_style.apply(prompt)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/generate/tp.py b/litgpt/generate/tp.py
index 9f10570..16bd5ac 100644
--- a/litgpt/generate/tp.py
+++ b/litgpt/generate/tp.py
@@ -104,6 +104,7 @@ def main(
     checkpoint_dir: Path,
     prompt: str = "What food do llamas eat?",
     *,
+    sys_prompt: Optional[str] = None,
     num_samples: int = 1,
     max_new_tokens: int = 50,
     top_k: Optional[int] = 50,
@@ -120,6 +121,7 @@ def main(
     Args:
         checkpoint_dir: The checkpoint directory to load.
         prompt: The prompt string to use for generating the samples.
+        sys_prompt: The system prompt to use for generating the samples.
         num_samples: The number of text samples to generate.
         max_new_tokens: The number of generation steps to take.
         top_k: The number of top most probable tokens to consider in the sampling process.
@@ -183,7 +185,7 @@ def main(
     prompt_style = (
         load_prompt_style(checkpoint_dir) if has_prompt_style(checkpoint_dir) else PromptStyle.from_config(config)
     )
-    prompt = prompt_style.apply(prompt)
+    prompt = prompt_style.apply(prompt, sys_prompt=sys_prompt)
     encoded = tokenizer.encode(prompt, device=fabric.device)
     prompt_length = encoded.size(0)
     max_returned_tokens = prompt_length + max_new_tokens
diff --git a/litgpt/prompts.py b/litgpt/prompts.py
index 17fda90..7297c6e 100644
--- a/litgpt/prompts.py
+++ b/litgpt/prompts.py
@@ -4,7 +4,7 @@ import re
 from abc import abstractmethod
 from json import dumps
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, List, Tuple, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type, Union
 
 import yaml
 
@@ -18,7 +18,7 @@ class PromptStyle:
     """Base interface for prompt styles."""
 
     @abstractmethod
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return prompt
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
@@ -34,7 +34,7 @@ class PromptStyle:
 
 
 class Default(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return prompt
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
@@ -42,47 +42,49 @@ class Default(PromptStyle):
 
 
 class Alpaca(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         if kwargs.get("input"):
-            return (
+            sys_prompt = sys_prompt or (
                 "Below is an instruction that describes a task, paired with an input that provides further context. "
                 "Write a response that appropriately completes the request.\n\n"
-                f"### Instruction:\n{prompt}\n\n### Input:\n{kwargs['input']}\n\n### Response:\n"
             )
-        return (
+            return f"{sys_prompt}### Instruction:\n{prompt}\n\n### Input:\n{kwargs['input']}\n\n### Response:\n"
+
+        sys_prompt = sys_prompt or (
             "Below is an instruction that describes a task. "
             "Write a response that appropriately completes the request.\n\n"
-            f"### Instruction:\n{prompt}\n\n### Response:\n"
         )
+        return f"{sys_prompt}### Instruction:\n{prompt}\n\n### Response:\n"
 
 
 class FLAN(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return (
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or (
             "Below is an instruction that describes a task. "
             "Write a response that appropriately completes the request.\n\n"
-            f"### Instruction:\n{prompt}\n\n### Response:\n"
         )
+        return f"{sys_prompt}### Instruction:\n{prompt}\n\n### Response:\n"
 
 
 class Longform(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return (
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or (
             "Below is an instruction that describes a task, paired with an input that provides further context. "
             "Write a response that appropriately completes the request.\n\n"
-            f"### Instruction:\n{prompt}\n\n### Response:\n"
         )
+        return f"{sys_prompt}### Instruction:\n{prompt}\n\n### Response:\n"
 
 
 class StableLMAlpha(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return (
-            "
# StableLM Tuned (Alpha version)\n- StableLM is a helpful and harmless open-source AI language"
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or (
+            "# StableLM Tuned (Alpha version)\n- StableLM is a helpful and harmless open-source AI language"
             " model developed by StabilityAI.\n- StableLM is excited to be able to help the user, but will refuse to do"
             " anything that could be considered harmful to the user.\n- StableLM is more than just an information"
             " source, StableLM is also able to write poetry, short stories, and make jokes.\n- StableLM will refuse to"
-            f" participate in anything that could harm a human.
{prompt}
"
+            " participate in anything that could harm a human."
         )
+        return f"
{sys_prompt}
{prompt}
"
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
         return (
@@ -94,12 +96,12 @@ class StableLMAlpha(PromptStyle):
 
 
 class StableLMZephyr(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"
\n{prompt}
\n
\n"
 
 
 class Falcon(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"{prompt}\nAnswer:"
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
@@ -113,7 +115,7 @@ class Falcon(PromptStyle):
 
 
 class Falcon3(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"
\n{prompt}
\n
\n"
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
@@ -124,7 +126,7 @@ class Falcon3(PromptStyle):
 
 
 class Llama2FunctionCalling(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         # Has to be before the llama config
         b_func, e_func = "<FUNCTIONS>", "</FUNCTIONS>\n\n"
         b_inst, e_inst = "
", "
"
@@ -139,7 +141,7 @@ class Llama2FunctionCalling(PromptStyle):
             "arguments": [{"name": "query", "type": "string", "description": "The search query string"}],
         }
 
-        system_prompt = (
+        system_prompt = sys_prompt or (
             "You are a helpful, respectful and honest assistant. Always answer as helpfully as"
             "possible. Your only response should be JSON formatted functions"
         )
@@ -151,22 +153,25 @@ class Llama2FunctionCalling(PromptStyle):
 
 
 class Llama2(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         b_inst, e_inst = "
", "
"
         b_sys, e_sys = "
\n", "\n
\n\n"
-        return (
-            f"{b_inst} {b_sys}You are a helpful, respectful and honest assistant. Always answer as helpfully as"
+        sys_prompt = sys_prompt or (
+            "You are a helpful, respectful and honest assistant. Always answer as helpfully as"
             " possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist,"
             " toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and"
             " positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why"
             " instead of answering something not correct. If you don't know the answer to a question, please don't"
-            f" share false information.{e_sys} {prompt} {e_inst} "
+            " share false information."
         )
+        return f"{b_inst} {b_sys}{sys_prompt}{e_sys} {prompt} {e_inst} "
 
 
 class Llama3(PromptStyle):
-    def apply(self, prompt: Union[str, List[Dict[str, str]]], **kwargs: str) -> str:
-        default_system_prompt = "You are a helpful assistant."
+    def apply(
+        self, prompt: Union[str, List[Dict[str, str]]], *, sys_prompt: Optional[str] = None, **kwargs: str
+    ) -> str:
+        default_system_prompt = sys_prompt or "You are a helpful assistant."
 
         # https://github.com/meta-llama/llama3/blob/359887376f0aaf30e433f23e25df858d8c2a9833/llama/tokenizer.py#L202-L229
         if isinstance(prompt, str):
@@ -216,8 +221,10 @@ class Llama3(PromptStyle):
 
 
 class R1Base(PromptStyle):
-    def apply(self, prompt: Union[str, List[Dict[str, str]]], **kwargs: str) -> str:
-        default_system_prompt = ""
+    def apply(
+        self, prompt: Union[str, List[Dict[str, str]]], *, sys_prompt: Optional[str] = None, **kwargs: str
+    ) -> str:
+        default_system_prompt = sys_prompt or ""
 
         bos_token = "<｜begin▁of▁sentence｜>"
         eos_token = ""
@@ -263,36 +270,35 @@ class R1Base(PromptStyle):
 
 
 class FreeWilly2(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return (
-            "### System:\nThis is a system prompt, please behave and help the user.\n\n"
-            "### User:\n"
-            f"{prompt}\n\n"
-            "### Assistant:\n"
-        )
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or "This is a system prompt, please behave and help the user."
+        return f"### System:\n{sys_prompt}\n\n### User:\n{prompt}\n\n### Assistant:\n"
 
 
 class Platypus(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"### Instruction:\n\n{prompt}\n\n### Response:\n"
 
 
 class StableCode(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"###Instruction\n{prompt}###Response\n"
 
 
 class CodeLlama(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         # for CodeLLama, we don't set a default system prompt, but it is supported:
         # https://huggingface.co/blog/codellama#conversational-instructions
         # Mistral does not: https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1#instruction-format
         b_inst, e_inst = "
", "
"
+        if sys_prompt:
+            b_sys, e_sys = "
\n", "\n
\n\n"
+            return f"{b_inst} {b_sys}{sys_prompt}{e_sys}{prompt} {e_inst}"
         return f"{b_inst} {prompt} {e_inst}"
 
 
 class Phi1(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"{prompt}\n\nAnswer:"
 
     def stop_tokens(self, tokenizer: "Tokenizer") -> Tuple[List[int], ...]:
@@ -307,38 +313,34 @@ class Phi1(PromptStyle):
 
 
 class Phi2(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"Instruct: {prompt}\nOutput:"
 
 
 class Phi3(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return f"
\nYou are a helpful assistant.
\n
\n{prompt}
\n
\n"
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or "You are a helpful assistant."
+        return f"
\n{sys_prompt}
\n
\n{prompt}
\n
\n"
 
 
 class Phi4(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"
user
{prompt}

assistant
"
 
 
 class TinyLlama(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return (
-            "
\n"
-            "You are a friendly chatbot who always gives helpful, detailed, and polite answers.
\n"
-            "
\n"
-            f"{prompt}
\n"
-            "
\n"
-        )
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or "You are a friendly chatbot who always gives helpful, detailed, and polite answers."
+        return f"
\n{sys_prompt}
\n
\n{prompt}
\n
\n"
 
 
 class Gemma(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
 
 
 class OLMo(PromptStyle):
-    def apply(self, prompt: str, **kwargs: str) -> str:
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
         return f"

\n{prompt}\n
\n"
 
 
@@ -346,8 +348,11 @@ class ChatML(PromptStyle):
     def __init__(self, system_message: str):
         self.system_message = system_message
 
-    def apply(self, prompt: str, **kwargs: str) -> str:
-        return f"
system\n{self.system_message}
\n
user\n{prompt}
\n
assistant\n"
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs: str) -> str:
+        sys_prompt = sys_prompt or self.system_message
+        return (
+            f"
system\n{sys_prompt}
\n
user\n{prompt}
\n
assistant\n"
+        )
 
 
 class Qwen2_5(ChatML):
diff --git a/litgpt/utils.py b/litgpt/utils.py
index e611d7f..cb588f9 100644
--- a/litgpt/utils.py
+++ b/litgpt/utils.py
@@ -63,13 +63,6 @@ def find_resume_path(resume: Union[bool, Literal["auto"], Path], out_dir: Path)
     return resume_path
 
 
-def find_multiple(n: int, k: int) -> int:
-    assert k > 0
-    if n % k == 0:
-        return n
-    return n + k - (n % k)
-
-
 def num_parameters(module: nn.Module, requires_grad: Optional[bool] = None) -> int:
     total = 0
     for p in module.parameters():


Test patch:
diff --git a/tests/data/test_base.py b/tests/data/test_base.py
index 52769d1..eb990f4 100644
--- a/tests/data/test_base.py
+++ b/tests/data/test_base.py
@@ -1,5 +1,7 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
+from typing import Optional
+
 import pytest
 import torch
 
@@ -12,7 +14,7 @@ from litgpt.prompts import PromptStyle
 @pytest.mark.parametrize("max_seq_length", [1000, 5, -1])
 def test_sft_dataset(max_seq_length, ignore_index, mask_prompt, mock_tokenizer):
     class Style(PromptStyle):
-        def apply(self, prompt, **kwargs):
+        def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs) -> str:
             return f"In: {prompt} Out:"
 
     i = ignore_index
diff --git a/tests/data/test_json.py b/tests/data/test_json.py
index 26345f4..e10141c 100644
--- a/tests/data/test_json.py
+++ b/tests/data/test_json.py
@@ -1,5 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import json
+from typing import Optional
 
 import pytest
 
@@ -10,7 +11,7 @@ from litgpt.prompts import PromptStyle
 @pytest.mark.parametrize("as_jsonl", [False, True])
 def test_json(as_jsonl, tmp_path, mock_tokenizer):
     class Style(PromptStyle):
-        def apply(self, prompt, **kwargs):
+        def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs) -> str:
             return f"X: {prompt} {kwargs['input']} Y:"
 
     json_path = tmp_path / ("data.jsonl" if as_jsonl else "data.json")
diff --git a/tests/test_config.py b/tests/test_config.py
index 925569e..a287357 100644
--- a/tests/test_config.py
+++ b/tests/test_config.py
@@ -5,6 +5,7 @@ import yaml
 
 import litgpt.config as config_module
 from litgpt import Config
+from litgpt.config import find_multiple
 
 
 def test_config():
@@ -103,3 +104,13 @@ def test_head_size(head_size):
     config = Config(head_size)
 
     assert config.head_size == head_size or config.n_embd // config.n_head
+
+
+def test_find_multiple():
+    assert find_multiple(17, 5) == 20
+    assert find_multiple(30, 7) == 35
+    assert find_multiple(10, 2) == 10
+    assert find_multiple(5, 10) == 10
+    assert find_multiple(50254, 128) == 50304
+    assert find_multiple(50254, 256) == 50432
+    assert find_multiple(50254, 512) == 50688
diff --git a/tests/test_prompts.py b/tests/test_prompts.py
index a211ffe..c882e6f 100644
--- a/tests/test_prompts.py
+++ b/tests/test_prompts.py
@@ -1,5 +1,8 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
+from typing import Optional
+
+import pytest
 import yaml
 
 import litgpt.config
@@ -8,6 +11,7 @@ from litgpt.prompts import (
     Alpaca,
     Default,
     Llama3,
+    Phi3,
     PromptStyle,
     has_prompt_style,
     load_prompt_style,
@@ -23,6 +27,26 @@ def test_default_prompt_style(mock_tokenizer):
     assert prompt_style.stop_tokens(mock_tokenizer) == ([mock_tokenizer.eos_id],)
 
 
+@pytest.mark.parametrize("sys_prompt", [None, "You are a helpful coding assistant."])
+def test_sys_prompt(mock_tokenizer, sys_prompt: Optional[str]):
+    prompt_style = Phi3()
+    prompt = "This is a test prompt."
+    default_sys_prompt = "You are a helpful assistant."
+    response = f"
\n{sys_prompt or default_sys_prompt}
\n
\n{prompt}
\n
\n"
+    assert prompt_style.apply(prompt, sys_prompt=sys_prompt) == response
+    assert prompt_style.stop_tokens(mock_tokenizer) == ([mock_tokenizer.eos_id],)
+
+
+@pytest.mark.parametrize("sys_prompt", [None, "You are a helpful coding assistant."])
+def test_sys_prompt_with_kwargs(mock_tokenizer, sys_prompt: Optional[str]):
+    prompt_style = Phi3()
+    prompt = "This is a test prompt."
+    default_sys_prompt = "You are a helpful assistant."
+    response = f"
\n{sys_prompt or default_sys_prompt}
\n
\n{prompt}
\n
\n"
+    assert prompt_style.apply(prompt, sys_prompt=sys_prompt, test=1) == response
+    assert prompt_style.stop_tokens(mock_tokenizer) == ([mock_tokenizer.eos_id],)
+
+
 def test_prompt_style_from_name():
     for style_name in prompt_styles:
         assert isinstance(PromptStyle.from_name(style_name), prompt_styles[style_name])
@@ -78,7 +102,7 @@ def test_apply_prompts():
 
 
 class CustomPromptStyle(PromptStyle):
-    def apply(self, prompt, **kwargs):
+    def apply(self, prompt: str, *, sys_prompt: Optional[str] = None, **kwargs) -> str:
         return prompt
 
 
diff --git a/tests/test_serve.py b/tests/test_serve.py
index 5f056a0..83d0568 100644
--- a/tests/test_serve.py
+++ b/tests/test_serve.py
@@ -160,27 +160,18 @@ def test_serve_with_openai_spec_missing_chat_template(tmp_path):
     def run_server():
         nonlocal process
         try:
-            process = subprocess.Popen(run_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
+            process = subprocess.Popen(run_command, stdout=None, stderr=None, text=True)
         except subprocess.TimeoutExpired:
             print("Server start-up timeout expired")
-        return None, None
 
     server_thread = threading.Thread(target=run_server)
     server_thread.start()
 
-    time.sleep(30)  # Give the server some time to start and raise the error
+    _wait_and_check_response()
 
-    try:
-        stdout = process.stdout.read().strip() if process.stdout else ""
-        stderr = process.stderr.read().strip() if process.stderr else ""
-        output = (stdout or "") + (stderr or "")
-        assert "ValueError: chat_template not found in tokenizer config file." in output, (
-            "Expected ValueError for missing chat_template not found."
-        )
-    finally:
-        if process:
-            kill_process_tree(process.pid)
-        server_thread.join()
+    if process:
+        kill_process_tree(process.pid)
+    server_thread.join()
 
 
 @_RunIf(min_cuda_gpus=1)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 13d8f01..4f41366 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -31,7 +31,6 @@ from litgpt.utils import (
     chunked_cross_entropy,
     copy_config_files,
     extend_checkpoint_dir,
-    find_multiple,
     find_resume_path,
     fix_and_load_json,
     incremental_save,
@@ -45,16 +44,6 @@ from litgpt.utils import (
 )
 
 
-def test_find_multiple():
-    assert find_multiple(17, 5) == 20
-    assert find_multiple(30, 7) == 35
-    assert find_multiple(10, 2) == 10
-    assert find_multiple(5, 10) == 10
-    assert find_multiple(50254, 128) == 50304
-    assert find_multiple(50254, 256) == 50432
-    assert find_multiple(50254, 512) == 50688
-
-
 # match fails on windows. why did they have to use backslashes?
 @_RunIf(skip_windows=True)
 def test_check_valid_checkpoint_dir(tmp_path):