Skip to content

update for MaterializedArray #489

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
davidkoski wants to merge 2 commits into
base: materialized-array-1
Choose a base branch
from
Open

update for MaterializedArray #489

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
443220e
update for MaterializedArray
davidkoski Jul 20, 2026
46f53c3
fix more warnings
davidkoski Jul 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 5 additions & 5 deletions Applications/LLMBasic/ChatModel.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,8 @@ private let generateParameters = GenerateParameters(temperature: 0.5)

enum State {
case idle
case loading(Task<ModelContainer, Error>)
case loaded(ModelContainer)
case loading(Task<ModelContext, Error>)
case loaded(ModelContext)
Comment on lines -27 to +28

@davidkoski davidkoski Jul 20, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Switching everything from ModelContainer -> ModelContext (now Sendable)

}

public var progress = 0.0
Expand All @@ -38,12 +38,12 @@ private let generateParameters = GenerateParameters(temperature: 0.5)

private var state = State.idle

public func model() async throws -> ModelContainer {
public func model() async throws -> ModelContext {
switch self.state {
case .idle:
let task = Task {
// download and report progress
try await #huggingFaceLoadModelContainer(
try await #huggingFaceLoadModel(
configuration: modelConfiguration
) { value in
Task { @MainActor in
Expand Down Expand Up @@ -79,7 +79,7 @@ private let generateParameters = GenerateParameters(temperature: 0.5)
task != nil
}

public init(model: ModelContainer) {
public init(model: ModelContext) {
self.session = ChatSession(
model,
instructions: instructions,
Expand Down
25 changes: 13 additions & 12 deletions Applications/LLMEval/ViewModels/LLMEvaluator.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ class LLMEvaluator {
enum LoadState {
case idle
case loading
case loaded(ModelContainer)
case loaded(ModelContext)
}

var loadState = LoadState.idle
Expand All @@ -81,7 +81,7 @@ class LLMEvaluator {
}

/// Load and return the model. Can be called multiple times; subsequent calls return the cached model.
func load() async throws -> ModelContainer {
func load() async throws -> ModelContext {
while true {
switch loadState {
case .idle:
Expand All @@ -91,13 +91,13 @@ class LLMEvaluator {
// Already loading, wait and retry
try await Task.sleep(for: .milliseconds(100))

case .loaded(let modelContainer):
return modelContainer
case .loaded(let model):
return model
}
}
}

private func performLoad() async throws -> ModelContainer {
private func performLoad() async throws -> ModelContext {
loadState = .loading
modelInfo = "Downloading \(modelName)..."
downloadProgress = 0.0
Expand Down Expand Up @@ -137,16 +137,16 @@ class LLMEvaluator {
downloadProgress = nil
totalSize = nil

let modelContainer = try await LLMModelFactory.shared.loadContainer(
let context = try await LLMModelFactory.shared.load(
from: resolved.modelDirectory,
using: #huggingFaceTokenizerLoader())

let numParams = await modelContainer.perform { $0.model.numParameters() }
let numParams = context.model.parameterCount

self.prompt = PresetPrompts.all[0].prompt
self.modelInfo = formatModelInfo(name: modelConfiguration.name, parameters: numParams)
loadState = .loaded(modelContainer)
return modelContainer
loadState = .loaded(context)
return context

} catch {
resetLoadingState()
Expand Down Expand Up @@ -240,18 +240,19 @@ class LLMEvaluator {
)

do {
let modelContainer = try await load()
let context = try await load()

// Capture parameters on MainActor before entering perform block
let parameters = generateParameters

// Seed random generator to ensure varied output each generation
MLXRandom.seed(UInt64(Date.timeIntervalSinceReferenceDate * 1000))

let lmInput = try await modelContainer.prepare(input: userInput)
let lmInput = try await context.processor.prepare(input: userInput)
let promptTokenCount = lmInput.text.tokens.size
let start = Date.timeIntervalSinceReferenceDate
let stream = try await modelContainer.generate(input: lmInput, parameters: parameters)
let stream = try MLXLMCommon.generate(
input: lmInput, parameters: parameters, context: context)

var iterator = stream.makeAsyncIterator()
if let first = await iterator.next() {
Expand Down
81 changes: 38 additions & 43 deletions Applications/LoRATrainingExample/ContentView.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -114,9 +114,9 @@ class LoRAEvaluator {
case failed(String)
}

enum ModelState: Sendable {
enum ModelState {
case idle
case loaded(ModelContainer)
case loaded(TrainableModelContext)
}

var state = State.idle
Expand All @@ -135,15 +135,15 @@ class LoRAEvaluator {
private let evaluateShowEvery = 8
private let maxTokens = 200

private func loadModel() async throws -> ModelContainer {
private func loadModel() async throws -> TrainableModelContext {
switch self.model {
case .idle:
let name = modelConfiguration.name
await MainActor.run {
progress = .init(title: "Loading \(name)", current: 0, limit: 1)
}

let modelContainer = try await #huggingFaceLoadModelContainer(
let context = try await #huggingFaceLoadTrainableModel(
configuration: modelConfiguration
) {
progress in
Expand All @@ -153,8 +153,8 @@ class LoRAEvaluator {
limit: 1.0)
}
}
self.model = .loaded(modelContainer)
return modelContainer
self.model = .loaded(context)
return context

case .loaded(let modelContainer):
return modelContainer
Expand Down Expand Up @@ -185,15 +185,13 @@ class LoRAEvaluator {
}

// load the model
let modelContainer = try await loadModel()
let context = try await loadModel()

// apply LoRA adapters and train
let _ = try await modelContainer.perform { context in
try LoRAContainer.from(
model: context.model,
configuration: LoRAConfiguration(numLayers: loraLayers)
)
}
let _ = try LoRAContainer.from(
model: context.model,
configuration: LoRAConfiguration(numLayers: loraLayers)
)

let train = try loadLoRAData(name: "train")
let valid = try loadLoRAData(name: "valid")
Expand All @@ -202,29 +200,27 @@ class LoRAEvaluator {
return
}

try await modelContainer.perform { context in
let optimizer = Adam(learningRate: learningRate)
try LoRATrain.train(
model: context.model, train: train, validate: valid, optimizer: optimizer,
tokenizer: context.tokenizer,
parameters: parameters
) { progress in
Task { @MainActor in
switch progress {
case .train(let i, _, _, _):
self.progress = .init(
title: "Train", current: Double(i), limit: Double(parameters.iterations)
)
case .validation:
output += "\n"
default:
break
}
output += progress.description + "\n"
let optimizer = Adam(learningRate: learningRate)
try LoRATrain.train(
model: context.model, train: train, validate: valid, optimizer: optimizer,
tokenizer: context.tokenizer,
parameters: parameters
) { progress in
Task { @MainActor in
switch progress {
case .train(let i, _, _, _):
self.progress = .init(
title: "Train", current: Double(i), limit: Double(parameters.iterations)
)
case .validation:
output += "\n"
default:
break
}

return .more
output += progress.description + "\n"
}

return .more
}

// done training, test
Expand All @@ -234,11 +230,9 @@ class LoRAEvaluator {
return
}

let loss = await modelContainer.perform { context in
LoRATrain.evaluate(
model: context.model, dataset: test,
tokenizer: context.tokenizer, batchSize: 1, batchCount: 0)
}
let loss = LoRATrain.evaluate(
model: context.model, dataset: test,
tokenizer: context.tokenizer, batchSize: 1, batchCount: 0)

self.progress = nil
self.output += "\n"
Expand All @@ -262,15 +256,16 @@ class LoRAEvaluator {

MLXRandom.seed(UInt64(Date.timeIntervalSinceReferenceDate * 1000))

let modelContainer = try await loadModel()
let context = try await loadModel()

// evaluate
let input = try await modelContainer.processor.prepare(input: .init(prompt: prompt))
let input = try await context.processor.prepare(input: .init(prompt: prompt))

let evaluationContext = ModelContext(context)
var count = 0
var output = ""
for try await item in try await modelContainer.generate(
input: input, parameters: generateParameters
for try await item in try generate(
input: input, parameters: generateParameters, context: evaluationContext
) {
switch item {
case .chunk(let string):
Expand Down
38 changes: 22 additions & 16 deletions Applications/MLXChatExample/Services/MLXService.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,8 +39,16 @@ class MLXService {
LMModel(name: "gemma3n:E4B", configuration: LLMRegistry.gemma3n_E4B_it_lm_4bit, type: .llm),
]

fileprivate final class ContextBox: Sendable {
let context: ModelContext

init(_ context: ModelContext) {
self.context = context
}
}

/// Cache to store loaded model containers to avoid reloading.
private let modelCache = NSCache<NSString, ModelContainer>()
private let modelCache = NSCache<NSString, ContextBox>()

/// Tracks the current model download progress.
/// Access this property to monitor model download status.
Expand All @@ -51,16 +59,16 @@ class MLXService {
/// - Parameter model: The model configuration to load
/// - Returns: A ModelContainer instance containing the loaded model
/// - Throws: Errors that might occur during model loading
private func load(model: LMModel) async throws -> ModelContainer {
private func load(model: LMModel) async throws -> ModelContext {
// Set GPU memory limit to prevent out of memory issues
Memory.cacheLimit = 20 * 1024 * 1024

// Return cached model if available to avoid reloading
if let container = modelCache.object(forKey: model.name as NSString) {
return container
if let box = modelCache.object(forKey: model.name as NSString) {
return box.context
} else {
// Select appropriate factory based on model type
let factory: ModelFactory =
let factory: any ModelFactory =
switch model.type {
case .llm:
LLMModelFactory.shared
Expand All @@ -72,7 +80,7 @@ class MLXService {
let loader = #huggingFaceTokenizerLoader()

// Load model and track download progress
let container = try await factory.loadContainer(
let context = try await factory.load(
from: downloader,
using: loader,
configuration: model.configuration
Expand All @@ -83,9 +91,9 @@ class MLXService {
}

// Cache the loaded model for future use
modelCache.setObject(container, forKey: model.name as NSString)
modelCache.setObject(.init(context), forKey: model.name as NSString)

return container
return context
}
}

Expand All @@ -97,7 +105,7 @@ class MLXService {
/// - Throws: Errors that might occur during generation
func generate(messages: [Message], model: LMModel) async throws -> AsyncStream<Generation> {
// Load or retrieve model from cache
let modelContainer = try await load(model: model)
let context = try await load(model: model)

// Exclude trailing empty assistant message so the chat template
// leaves the assistant turn open for generation (matching ChatSession behavior)
Expand Down Expand Up @@ -131,13 +139,11 @@ class MLXService {
chat: chat, processing: .init(resize: .init(width: 1024, height: 1024)))

// Generate response using the model
return try await modelContainer.perform { context in
let lmInput = try await context.processor.prepare(input: userInput)
// Set temperature for response randomness (0.7 provides good balance)
let parameters = GenerateParameters(temperature: 0.7)
let lmInput = try await context.processor.prepare(input: userInput)
// Set temperature for response randomness (0.7 provides good balance)
let parameters = GenerateParameters(temperature: 0.7)

return try MLXLMCommon.generate(
input: lmInput, parameters: parameters, context: context)
}
return try MLXLMCommon.generate(
input: lmInput, parameters: parameters, context: context)
}
}
Loading
Loading