Finetune demo

crates/piston-core/src/gpu/buffer_allocator/allocator.rs

@@ -7,7 +7,7 @@ use crate::{
         PooledGPUBuffer, TensorUsageRecords, UNIFORM_ALIGN, WgpuDevice,
     },
 };
-use crate::{HashMap, LazyOp};
+use crate::{HashMap, HashSet, LazyOp};
 use parking_lot::RwLock;
 use std::num::NonZero;
 use std::{borrow::Cow, collections::BTreeMap};
@@ -281,6 +281,7 @@ impl BufferAllocator {
         output_tensors: &BTreeMap<TensorId, &OpTensor>,
         assignments: &mut HashMap<TensorId, PooledGPUBuffer>,
         gpu_compile_keys: &HashMap<TensorId, GpuCompileKey>,
+        shareable_ids: &HashSet<TensorId>,
         use_shared_buffers: bool,
         device: &WgpuDevice,
     ) -> Result<(), DeviceError> {
@@ -289,9 +290,16 @@ impl BufferAllocator {
         let mut shared_objects: Vec<PooledGPUBuffer> = Vec::with_capacity(records.0.len());
 
         for record in records.0.iter() {
+            let is_output = output_tensors.get(&record.last_consumer_id).is_some();
+            let is_shareable_id = record
+                .id
+                .map(|id| shareable_ids.contains(&id))
+                .unwrap_or(false);
+
             let should_be_shared = use_shared_buffers
-                && !(record.requires_grad.unwrap_or(false)
-                    || output_tensors.get(&record.last_consumer_id).is_some());
+                && is_shareable_id
+                && !is_output
+                && !record.requires_grad.unwrap_or(false);
 
             let mut best_obj = None;
 
@@ -370,6 +378,7 @@ impl BufferAllocator {
         execution_order: &[&OpTensor],
         output_tensors: &BTreeMap<TensorId, &OpTensor>,
         gpu_compile_keys: &HashMap<TensorId, GpuCompileKey>,
+        shareable_ids: &HashSet<TensorId>,
         use_shared_buffers: bool,
         device: &WgpuDevice,
     ) -> Result<HashMap<TensorId, PooledGPUBuffer>, DeviceError> {
@@ -395,6 +404,7 @@ impl BufferAllocator {
             output_tensors,
             &mut assignments,
             gpu_compile_keys,
+            shareable_ids,
             use_shared_buffers,
             device,
         )?;

crates/piston-core/src/gpu/buffer_allocator/lazy_graph_executor.rs

@@ -1,7 +1,9 @@
 use crate::{
     Compiled, CpuUniform, DebugSelection, Executable, ExecutionError, ExecutionResult, GPUBuffer,
     HashMap, HashSet, Hasher as HasherType, Inner, LazyOp, StepLog, StepLogConfig, Storage,
-    TensorError, WgpuDevice, reset_scope_context,
+    TensorError, WgpuDevice,
+    gpu::{TensorOpDesc, trace_sink},
+    reset_scope_context,
 };
 #[cfg(feature = "debug")]
 use crate::{DebugTensor, Device, DeviceStorage};
@@ -392,7 +394,34 @@ impl LazyGraphExecutor {
             }
         }
 
-        log::debug!("Post-order hash: {:?}", hasher.finish());
+        // Determine which tensors are safe candidates for shared-object buffers.
+        //
+        // A tensor is considered shareable if:
+        // - It does not have extra external strong references beyond what we expect from the
+        //   execution graph and (optionally) owned_tensors.
+        // - It is not a leaf / parameter (requires_grad).
+        // - It does not retain its gradient.
+        //
+        // These tensors are good candidates for hosting shared-object buffers whose contents may
+        // be reused for other logical tensors once their lifetimes (based on execution_order)
+        // have ended.
+        let mut shareable_ids =
+            HashSet::with_capacity_and_hasher(post_order.len(), Default::default());
+        for t in &post_order {
+            let id = t.id();
+            let expected_strong = owned_tensors
+                .as_ref()
+                .and_then(|ot| ot.contains(&id).then_some(2))
+                .unwrap_or(1);
+            let strong = t.strong_count();
+
+            let externally_pinned =
+                strong > expected_strong || t.requires_grad() || t.retains_grad();
+
+            if !externally_pinned {
+                shareable_ids.insert(id);
+            }
+        }
 
         let output_tensors = tensors
             .iter()
@@ -483,6 +512,7 @@ impl LazyGraphExecutor {
                 &post_order,
                 &output_tensors,
                 &compile_keys,
+                &shareable_ids,
                 self.shared_object_allocation_enabled,
                 gpu_device,
             )?)

crates/piston-core/src/gpu/device.rs

@@ -304,13 +304,15 @@ impl WgpuDevice {
         execution_order: &[&OpTensor],
         output_tensors: &BTreeMap<TensorId, &OpTensor>,
         gpu_compile_keys: &HashMap<TensorId, GpuCompileKey>,
+        shareable_ids: &crate::HashSet<TensorId>,
         use_shared_buffers: bool,
         device: &WgpuDevice,
     ) -> Result<HashMap<TensorId, PooledGPUBuffer>, DeviceError> {
         self.buffer_allocator.allocate_cfg(
             execution_order,
             output_tensors,
             gpu_compile_keys,
+            shareable_ids,
             use_shared_buffers,
             device,
         )

examples/finetuning/.editorconfig

@@ -0,0 +1,23 @@
+root = true
+
+[*]
+end_of_line = lf
+insert_final_newline = true
+indent_style = tab
+indent_size = tab
+tab_width = 2
+charset = utf-8
+trim_trailing_whitespace = true
+
+[*.py]
+indent_style = space
+indent_size = 4
+tab_width = 4
+
+[*.ipynb]
+indent_style = space
+indent_size = 4
+tab_width = 4
+
+[package.json]
+indent_style = space

examples/finetuning/.gitignore

@@ -0,0 +1,28 @@
+node_modules
+
+# Output
+.output
+.vercel
+.netlify
+.wrangler
+/.svelte-kit
+/build
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Env
+*.local
+.env*
+!.env.production
+!.env.development
+!.env.example
+!.env.test
+
+# Vite
+vite.config.js.timestamp-*
+vite.config.ts.timestamp-*
+
+static/tokenizer
+static/tokenized

examples/finetuning/.npmrc

@@ -0,0 +1,2 @@
+engine-strict=true
+use-node-version=22.17.1

examples/finetuning/.prettierignore

@@ -0,0 +1,4 @@
+# Package Managers
+package-lock.json
+pnpm-lock.yaml
+yarn.lock

examples/finetuning/.prettierrc

@@ -0,0 +1,18 @@
+{
+	"useTabs": true,
+	"singleQuote": true,
+	"trailingComma": "none",
+	"printWidth": 100,
+	"tabWidth": 2,
+	"plugins": [
+		"prettier-plugin-svelte"
+	],
+	"overrides": [
+		{
+			"files": "*.svelte",
+			"options": {
+				"parser": "svelte"
+			}
+		}
+	]
+}

examples/finetuning/README.md

@@ -0,0 +1 @@
+This is the source code for the finetuning demo at [finetune.sequence.toys](https://finetune.sequence.toys). It is a static Svelte app.

examples/finetuning/eslint.config.js

@@ -0,0 +1,81 @@
+import { includeIgnoreFile } from '@eslint/compat';
+import js from '@eslint/js';
+import prettier from 'eslint-config-prettier';
+import perfectionist from 'eslint-plugin-perfectionist';
+import svelte from 'eslint-plugin-svelte';
+import { defineConfig } from 'eslint/config';
+import globals from 'globals';
+import { fileURLToPath } from 'node:url';
+import ts from 'typescript-eslint';
+const gitignorePath = fileURLToPath(new URL('./.gitignore', import.meta.url));
+const tsconfigRootDir = fileURLToPath(new URL('.', import.meta.url));
+
+export default defineConfig(
+	includeIgnoreFile(gitignorePath),
+	js.configs.recommended,
+	...ts.configs.recommended,
+	...svelte.configs['flat/recommended'],
+	...svelte.configs['flat/prettier'],
+	{
+		languageOptions: {
+			globals: {
+				...globals.browser,
+				...globals.node
+			}
+		}
+	},
+	{
+		files: ['**/*.svelte', '**/*.svelte.ts', '**/*.ts'],
+
+		languageOptions: {
+			parserOptions: {
+				parser: ts.parser,
+				// Ensure the TypeScript parser knows which tsconfig to use in this package
+				tsconfigRootDir,
+				projectService: true
+			}
+		}
+	},
+	{
+		plugins: {
+			prettier,
+			perfectionist
+		},
+		rules: {
+			'no-unused-vars': 'off',
+			'@typescript-eslint/ban-types': 'off',
+			'@typescript-eslint/no-explicit-any': 'warn',
+			'@typescript-eslint/no-unused-vars': [
+				'warn',
+				{
+					argsIgnorePattern: '^_',
+					varsIgnorePattern: '^_',
+					caughtErrorsIgnorePattern: '^_'
+				}
+			],
+			'@typescript-eslint/no-var-requires': 'off',
+			'perfectionist/sort-imports': [
+				'error',
+				{
+					type: 'natural',
+					order: 'asc',
+					groups: [
+						'type',
+						['builtin', 'external'],
+						'internal-type',
+						'internal',
+						['parent-type', 'sibling-type', 'index-type'],
+						['parent', 'sibling', 'index'],
+						'side-effect',
+						'style',
+						'object',
+						'unknown'
+					]
+				}
+			],
+			'perfectionist/sort-exports': ['error', { type: 'natural' }],
+			'perfectionist/sort-named-exports': ['error', { type: 'natural' }],
+			'perfectionist/sort-named-imports': ['error', { type: 'natural' }]
+		}
+	}
+);

examples/finetuning/package.json

@@ -0,0 +1,63 @@
+{
+  "name": "piston-finetune-toy",
+  "private": true,
+  "version": "0.0.1",
+  "type": "module",
+  "scripts": {
+    "dev": "vite dev",
+    "build": "vite build",
+    "preview": "vite preview",
+    "prepare": "svelte-kit sync || echo ''",
+    "check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json",
+    "check:watch": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json --watch",
+    "lint": "eslint . && prettier --check .",
+    "format": "prettier --write ."
+  },
+  "dependencies": {
+    "@codemirror/autocomplete": "^6.19.1",
+    "@codemirror/language": "^6.11.3",
+    "@codemirror/lint": "^6.9.1",
+    "@codemirror/state": "^6.5.2",
+    "@codemirror/view": "^6.38.6",
+    "@huggingface/jinja": "^0.5.1",
+    "@lezer/highlight": "^1.2.3",
+    "@lucide/svelte": "^0.554.0",
+    "@piston-ml/piston-web": "workspace:^",
+    "@types/katex": "^0.16.7",
+    "codemirror": "^6.0.2",
+    "echarts": "^6.0.0",
+    "example-common": "workspace:*",
+    "katex": "^0.16.23",
+    "maxrects-packer": "^2.7.3",
+    "random-js": "^2.1.0",
+    "svelte-portal": "^2.2.1",
+    "unique-names-generator": "^4.7.1"
+  },
+  "devDependencies": {
+    "@eslint/compat": "^1.4.0",
+    "@eslint/js": "^9.37.0",
+    "@sveltejs/adapter-static": "^3.0.10",
+    "@sveltejs/kit": "^2.46.4",
+    "@sveltejs/vite-plugin-svelte": "^6.2.1",
+    "@tailwindcss/vite": "^4.1.14",
+    "@types/glob": "^9.0.0",
+    "@webgpu/types": "^0.1.65",
+    "eslint": "^9.37.0",
+    "eslint-config-prettier": "^10.1.8",
+    "eslint-plugin-perfectionist": "^4.15.1",
+    "eslint-plugin-svelte": "^3.12.4",
+    "glob": "^11.0.3",
+    "globals": "^16.4.0",
+    "prettier": "^3.6.2",
+    "prettier-plugin-svelte": "^3.4.0",
+    "rollup": "^4.52.4",
+    "sirv": "^2.0.4",
+    "svelte": "workspace:^",
+    "svelte-check": "^4.3.3",
+    "tailwindcss": "^4.1.14",
+    "typescript": "^5.9.3",
+    "typescript-eslint": "^8.46.0",
+    "vite": "^7.1.9",
+    "vite-plugin-wasm": "^3.5.0"
+  }
+}

examples/finetuning/src/app.css

@@ -0,0 +1,25 @@
+@import 'tailwindcss' source('./');
+@import 'example-common/theme.css';
+
+@source "../../../packages/example-common/src/**/*.{svelte,js,ts}";
+
+html {
+	/* Base font size for fine pointers (desktop/mouse) */
+	font-size: 16px;
+	overscroll-behavior: none;
+	@apply h-full;
+	@apply w-full;
+}
+
+@media (pointer: coarse) {
+	html {
+		font-size: 18px;
+	}
+}
+
+body {
+	@apply h-full;
+	@apply w-full;
+	@apply text-base;
+	overscroll-behavior: none;
+}

examples/finetuning/src/app.d.ts

@@ -0,0 +1,14 @@
+// See https://svelte.dev/docs/kit/types#app.d.ts
+// for information about these interfaces
+declare global {
+	const __COMMIT_HASH__: string;
+	namespace App {
+		// interface Error {}
+		// interface Locals {}
+		// interface PageData {}
+		// interface PageState {}
+		// interface Platform {}
+	}
+}
+
+export {};