[OpenAI] Support seed for reproducible generation (mlc-ai#335)

CharlieFRuan · web-flow · commit 9ffc6833bcbb · 2024-03-15T10:24:53.000-04:00
We support `seed` in `ChatCompletionRequest` so requests' results are
reproducible.

As stated in the docstring in
`src/openai_api_protocols/chat_completion.ts`, seeding is done at a
request level, rather than a choice level. So if a request with `n &gt; 1`
is seeded, the choices would still have different results. But if two
requests, both with `n &gt; 1` and share the same seed, would generate
identical results, across all choices. This is demonstrated in
`examples/openai-api/src/seed.ts`, where we rigorously compare the
strings generated.

Implementation wise, this is achieved with a customized implementation
of linear congruential generator in TVMjs's runtime, since JS's
`Math.random()` does not support seeding.
diff --git a/examples/openai-api/README.md b/examples/openai-api/README.md
@@ -0,0 +1,19 @@
+### OpenAI API Demos
+
+Run `npm install` first, followed by `npm start`.
+
+To run different scripts, you can modify `package.json` from the default 
+```json
+"scripts": {
+    "start": "parcel src/openai_api.html  --port 8888",
+    "build": "parcel build src/openai_api.html --dist-dir lib"
+},
+```
+
+to, say
+```json
+"scripts": {
+    "start": "parcel src/seed.html  --port 8888",
+    "build": "parcel build src/seed.html --dist-dir lib"
+},
+```
diff --git a/examples/openai-api/src/seed.html b/examples/openai-api/src/seed.html
@@ -0,0 +1,17 @@
+<!DOCTYPE html>
+<html>
+<script>
+    webLLMGlobal = {}
+</script>
+
+<body>
+    <h2>WebLLM OpenAI-like API Test Page</h2>
+    Open console to see output.
+    We make two generations with same seed, we should expect them to be the same.
+    </br>
+    </br>
+    <label id="init-label"> </label>
+
+    <script type="module" src="./seed.ts"></script>
+
+</html>
diff --git a/examples/openai-api/src/seed.ts b/examples/openai-api/src/seed.ts
@@ -0,0 +1,58 @@
+import * as webllm from "@mlc-ai/web-llm";
+
+function setLabel(id: string, text: string) {
+    const label = document.getElementById(id);
+    if (label == null) {
+        throw Error("Cannot find label " + id);
+    }
+    label.innerText = text;
+}
+
+/**
+ * We domnstrate the effect of seeding. The prompt is about writing a poem and we use a high
+ * `temperature`, making the sampling distribution supposedly more random. However, we demonstrate
+ * that with seeding, we should see the exact same result being generated across two trials.
+ * With `n > 1`, all choices should also be exactly the same.
+ */
+async function demonstrateSeed() {
+    const chat: webllm.ChatInterface = new webllm.ChatModule();
+
+    chat.setInitProgressCallback((report: webllm.InitProgressReport) => {
+        setLabel("init-label", report.text);
+    });
+
+    await chat.reload("Llama-2-7b-chat-hf-q4f32_1");
+
+    const request: webllm.ChatCompletionRequest = {
+        stream: false,  // works with streaming as well
+        messages: [
+            { "role": "user", "content": "Write a creative Haiku about Pittsburgh" }
+        ],
+        n: 3,
+        temperature: 1.2,  // high temperature gives much more random results
+        seed: 42,
+        max_gen_len: 128,  // To save time; enough to demonstrate the effect
+    };
+
+    const reply0 = await chat.chatCompletion(request);
+    console.log(reply0);
+    console.log("First reply's last choice:\n" + await chat.getMessage());
+
+    const reply1 = await chat.chatCompletion(request);
+    console.log(reply1);
+    console.log("Second reply's last choice:\n" + await chat.getMessage());
+
+    // Rigorously check the generation results of each choice for the two requests
+    for (const choice0 of reply0.choices) {
+        const id = choice0.index;
+        const choice1 = reply1.choices[id];
+        if (choice0.message.content !== choice1.message.content) {
+            throw Error("Chocie " + id + " of the two generations are different despite seeding");
+        }
+    }
+
+    console.log(await chat.runtimeStatsText());
+}
+
+// Run one of the functions
+demonstrateSeed();
diff --git a/src/chat_module.ts b/src/chat_module.ts
@@ -208,6 +208,9 @@ export class ChatModule implements ChatInterface {
     genConfig: GenerationConfig
   ): AsyncGenerator<ChatCompletionChunk, void, void> {
     postInitAndCheckGenerationConfigValues(genConfig);
+    if (request.seed !== null && request.seed !== undefined) {
+      this.getPipeline().setSeed(request.seed);
+    }
     if (!request.stateful) {
       await this.resetChat();
     }
@@ -255,6 +258,11 @@ export class ChatModule implements ChatInterface {
       yield await _getChunk(this);
     }
 
+    // Reset seed -- we do not want this seed to affect future requests
+    if (request.seed !== null && request.seed !== undefined) {
+      this.getPipeline().setSeed(Date.now());
+    }
+
     const lastChunk: ChatCompletionChunk = {
       id: id,
       choices: [{
@@ -270,6 +278,15 @@ export class ChatModule implements ChatInterface {
     yield lastChunk;
   }
 
+  /**
+   * Completes a single ChatCompletionRequest.
+   * 
+   * @param request A OpenAI-style ChatCompletion request.
+   * 
+   * @note For each choice (i.e. `n`), a request is defined by a single `prefill()` and mulitple
+   * `decode()`. This is important as it determines the behavior of various fields including
+   * `stateful` and `seed`.
+   */
   async chatCompletion(
     request: ChatCompletionRequestNonStreaming
   ): Promise<ChatCompletion>;
@@ -304,6 +321,10 @@ export class ChatModule implements ChatInterface {
       return this.chatCompletionAsyncChunkGenerator(request, genConfig);
     }
 
+    if (request.seed !== null && request.seed !== undefined) {
+      this.getPipeline().setSeed(request.seed);
+    }
+
     // 2. If request is non-streaming, directly reuse `generate()`
     const n = request.n ? request.n : 1;
     const choices: Array<ChatCompletion.Choice> = [];
@@ -354,6 +375,11 @@ export class ChatModule implements ChatInterface {
         total_tokens: completion_tokens + prompt_tokens,
       } as CompletionUsage,
     }
+
+    // Reset seed -- we do not want this seed to affect future requests
+    if (request.seed !== null && request.seed !== undefined) {
+      this.getPipeline().setSeed(Date.now());
+    }
     return response;
   }
 
diff --git a/src/llm_chat.ts b/src/llm_chat.ts
@@ -343,7 +343,14 @@ export class LLMChatPipeline {
     )
   }
 
-  // Getters and writters for this.conversation.
+  /**
+   * Set the seed for the RNG `this.tvm.rng`.
+   */
+  setSeed(seed: number): void {
+    this.tvm.setSeed(seed);
+  }
+
+  // Getters and setters for this.conversation.
   /**
    * Overrides the system prompt.
    */
diff --git a/src/openai_api_protocols/chat_completion.ts b/src/openai_api_protocols/chat_completion.ts
@@ -134,6 +134,16 @@ export interface ChatCompletionRequestBase {
      */
     top_logprobs?: number | null;
 
+    /**
+     * If specified, our system will make a best effort to sample deterministically, such that
+     * repeated requests with the same `seed` and parameters should return the same result.
+     * 
+     * @note Seeding is done on a request-level rather than choice-level. That is, if `n > 1`, you
+     * would still get different content for each `Chocie`. But if two requests with `n = 2` are
+     * processed with the same seed, the two results should be the same (two choices are different).
+     */
+    seed?: number | null;
+
     //////////////// BELOW FIELDS NOT SUPPORTED YET ////////////////
 
     /**
@@ -143,14 +153,6 @@ export interface ChatCompletionRequestBase {
      */
     model?: string | null;
 
-    /**
-     * If specified, our system will make a best effort to sample deterministically, such that
-     * repeated requests with the same `seed` and parameters should return the same result.
-     * 
-     * @note Not supported yet.
-     */
-    seed?: number | null;
-
     /**
      * Controls which (if any) function is called by the model. `none` means the model
      * will not call a function and instead generates a message. `auto` means the model
@@ -306,7 +308,6 @@ export const ChatCompletionRequestUnsupportedFields: Array<string> = [
     "tool_choice",
     "tools",
     "response_format",
-    "seed",
 ];
 
 export function postInitAndCheckFields(request: ChatCompletionRequest): void {
@@ -363,6 +364,13 @@ export function postInitAndCheckFields(request: ChatCompletionRequest): void {
     if (request.stateful && request.n && request.n > 1) {
         throw new Error("If the request is stateful, `n` cannot be > 1.");
     }
+
+    // 6. Seed should be an integer
+    if (request.seed !== undefined && request.seed !== null) {
+        if (!Number.isInteger(request.seed)) {
+            throw new Error("`seed` should be an integer, but got " + request.seed);
+        }
+    }
 }
 
 //////////////// BELOW ARE INTERFACES THAT SUPPORT THE ONES ABOVE ////////////////
diff --git a/tests/openai_chat_completion.test.ts b/tests/openai_chat_completion.test.ts
@@ -68,6 +68,19 @@ describe('Check chat completion unsupported requests', () => {
         }).toThrow("If the request is stateful, `n` cannot be > 1.");
     });
 
+    test('Non-integer seed', () => {
+        expect(() => {
+            const request: ChatCompletionRequest = {
+                messages: [
+                    { role: "user", content: "Hello! " },
+                ],
+                max_gen_len: 10,
+                seed: 42.2,  // Note that Number.isInteger(42.0) is true 
+            };
+            postInitAndCheckFields(request)
+        }).toThrow("`seed` should be an integer, but got");
+    });
+
     // Remove when we support image input (e.g. LlaVA model)
     test('Image input is unsupported', () => {
         expect(() => {
@@ -128,6 +141,14 @@ describe('Supported requests', () => {
             temperature: 1.5,
             max_gen_len: 25,
             frequency_penalty: 0.2,
+            seed: 42,
+            logprobs: true,
+            top_logprobs: 2,
+            logit_bias: {
+                "13813": -100,
+                "10319": 5,
+                "7660": 5,
+            },
         };
         postInitAndCheckFields(request)
     });
