adding FIM finetuned model hosted on fireworks #4245

hitesh-1997 · 2024-05-21T22:08:23Z

Context

This PR adds the fine-tuning model setup, trained on context aware FIM examples on top of bigcode/the-stack. More details about the training dataset/ modeling/ offline evaluation etc can be found in this RFC.
We wanted to iterate faster on the modeling changes, i.e. don't want to replace the model string in the client every time we change a model in the backend - So, this PR adds four variant feature flags and control feature flag and the variant based feature-flag names will be send to cody-gateway as is, where the routing to the actual fine-tuned model happens.
The feature flag cody-autocomplete-fim-finetuned-model-base-flag added in the setup is treated as the base feature flag and the traffic on this feature flag will decide the traffic of the whole experiment. The traffic on this feature will be divided further into the variant & control feature flags.

Test plan

Traffic split sanity:

I have tested the approach of feature flags introduced in this PR on random generated 100k userids offline and verified that it produces near equal traffic split among the variants and control. For anyone want to reproduce this results, can be done via this script.

New modeling changes:

Local testing

valerybugakov

LGTM!

valerybugakov · 2024-05-22T04:24:38Z

vscode/src/completions/providers/fireworks.ts

@@ -83,8 +95,10 @@ const MODEL_MAP = {
    'llama-code-13b': 'fireworks/accounts/fireworks/models/llama-v2-13b-code',

    // Fine-tuned model mapping
-    'fireworks-completions-fine-tuned':
-        'fireworks/accounts/sourcegraph/models/codecompletion-mixtral-rust-152k-005e',
+    'fim-fine-tuned-model-variant-1': FIREWORKS_FIM_FINE_TUNED_MODEL_1,


Since we have constants already, we can:

Suggested change

'fim-fine-tuned-model-variant-1': FIREWORKS_FIM_FINE_TUNED_MODEL_1,

[FIREWORKS_FIM_FINE_TUNED_MODEL_1]: FIREWORKS_FIM_FINE_TUNED_MODEL_1,

valerybugakov · 2024-05-22T04:31:33Z

vscode/src/completions/providers/create-provider.ts

+async function resolveFinetunedModelProviderFromFeatureFlags(): Promise<{
+    provider: string
+    model?: FireworksOptions['model'] | AnthropicOptions['model']
+} | null> {


To keep return types in sync automatically:

Suggested change

async function resolveFinetunedModelProviderFromFeatureFlags(): Promise<{

provider: string

model?: FireworksOptions['model'] | AnthropicOptions['model']

} | null> {

async function resolveFinetunedModelProviderFromFeatureFlags(): ReturnType<

typeof resolveDefaultProviderFromVSCodeConfigOrFeatureFlags

> {

wow nice, didn't about this.
Changed as per suggestion

valerybugakov · 2024-05-22T04:31:42Z

vscode/src/completions/providers/create-provider.ts

+
+    if (finetunedFIMModelExperiment) {
+        // The traffic in this feature flag is interpreted as a traffic allocated to the fine-tuned experiment.
+        return await resolveFinetunedModelProviderFromFeatureFlags()


Suggested change

return await resolveFinetunedModelProviderFromFeatureFlags()

return resolveFinetunedModelProviderFromFeatureFlags()

janhartman · 2024-05-22T07:44:59Z

vscode/src/completions/providers/create-provider.ts

+async function resolveFinetunedModelProviderFromFeatureFlags(): Promise<{
+    provider: string
+    model?: FireworksOptions['model'] | AnthropicOptions['model']
+} | null> {


rafax · 2024-05-22T08:49:48Z

lib/shared/src/experimentation/FeatureFlagProvider.ts

+    // Enable various feature flags to experiment with FIM trained fine-tuned models via Fireworks
+    CodyAutocompleteFIMFineTunedModelBaseFeatureFlag = 'cody-autocomplete-fim-finetuned-model-base-flag',
+    CodyAutocompleteFIMFineTunedModelControl = 'cody-autocomplete-fim-finetuned-model-control',
+    CodyAutocompleteFIMFineTunedModelVariant1 = 'cody-autocomplete-fim-finetuned-model-variant1',


Can we standardize the naming (prefix, variantX, spelling of fine-tuned) between this and internal/completions/client/fireworks/fireworks.go so we don't have to redeploy due to typos?

totally makes sense, changed the string names so that they are inline with names at other places.
I have followed the following convention here: _<variant_name>, eg: cody-autocomplete-fim-fine-tuned-model-variant-1 Does this look alright.

rafax · 2024-05-22T08:50:14Z

vscode/src/completions/providers/fireworks.ts

@@ -69,6 +68,19 @@ const PROVIDER_IDENTIFIER = 'fireworks'
 const EOT_STARCODER = '<|endoftext|>'
 const EOT_LLAMA_CODE = ' <EOT>'

+// Fireworks hosted model identifier strings


Link to the actual models backing those "virtual" identifiers?

I think we can add a link to sourcegraph repo where the acutal mapping is defined ?

philipp-spiess · 2024-05-22T11:10:53Z

vscode/src/completions/providers/fireworks.ts

+        case FIREWORKS_FIM_FINE_TUNED_MODEL_4: {
+            // We use llama3 8b and mixtral 8x7b variants for the fine-tuning model which support 8_192, 32_768 tokens respectively.
+            // Take a buffer of 1000 tokens
+            return 7192


Note that this is going to increase the number of context we send by a lot which might have an impact on the comparison and the latency

Thanks for noticing this philipp. Earlier I did a load test on different context token length from a GCP VM in us-central-1a and got ~100ms delta for p75, so I went ahead with this context length. But I realise the user is going to query this from their machine and bigger context can lead to increased latency when client is local.
changing this back to 2048 same as others.

I wouldn't be too concerned about the user/cpu overheard. 100ms for a 4x in tokens could be great tread off, though. The problem we ran into the last time was that it reduced throughput a lot as well. Not sure how this affects the setup.

adding FIM finetuned model hosted on fireworks

32a5cb7

hitesh-1997 requested review from valerybugakov, philipp-spiess and a team May 21, 2024 22:08

valerybugakov approved these changes May 22, 2024

View reviewed changes

hitesh-1997 mentioned this pull request May 22, 2024

adding finetuned model setup sourcegraph/sourcegraph#62838

Merged

hitesh-1997 requested a review from rafax May 22, 2024 07:24

janhartman approved these changes May 22, 2024

View reviewed changes

rafax reviewed May 22, 2024

View reviewed changes

philipp-spiess reviewed May 22, 2024

View reviewed changes

hitesh-1997 added 4 commits May 22, 2024 19:36

adding fix as per PR comments

7c71188

add languageId to the fireworks payload

d9f348c

run fine tuning experiment only for vscode client

cae9ec2

change context window

bdb8cfb

hitesh-1997 merged commit 4b9fb4c into main May 23, 2024
19 checks passed

hitesh-1997 deleted the hitesh/finetuned-fim-model-integration branch May 23, 2024 11:35

hitesh-1997 added a commit that referenced this pull request May 23, 2024

adding FIM finetuned model hosted on fireworks (#4245)

2150618

hitesh-1997 mentioned this pull request May 23, 2024

Hitesh/1.18.2 patch release #4273

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

adding FIM finetuned model hosted on fireworks #4245

adding FIM finetuned model hosted on fireworks #4245

hitesh-1997 commented May 21, 2024 •

edited

valerybugakov left a comment

valerybugakov May 22, 2024

hitesh-1997 May 22, 2024

valerybugakov May 22, 2024

janhartman May 22, 2024

hitesh-1997 May 22, 2024

valerybugakov May 22, 2024

hitesh-1997 May 22, 2024

janhartman May 22, 2024

rafax May 22, 2024

hitesh-1997 May 22, 2024

rafax May 22, 2024

hitesh-1997 May 22, 2024

philipp-spiess May 22, 2024

hitesh-1997 May 22, 2024

philipp-spiess May 22, 2024

	'fim-fine-tuned-model-variant-1': FIREWORKS_FIM_FINE_TUNED_MODEL_1,
	[FIREWORKS_FIM_FINE_TUNED_MODEL_1]: FIREWORKS_FIM_FINE_TUNED_MODEL_1,

	return await resolveFinetunedModelProviderFromFeatureFlags()
	return resolveFinetunedModelProviderFromFeatureFlags()

adding FIM finetuned model hosted on fireworks #4245

adding FIM finetuned model hosted on fireworks #4245

Conversation

hitesh-1997 commented May 21, 2024 • edited

Context

Test plan

Traffic split sanity:

New modeling changes:

valerybugakov left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

hitesh-1997 commented May 21, 2024 •

edited