Model Download Slow and Unstable

[CodeWithShreyans]

Using the Hub loadContainer function, the model download bandwidth take a sharp fall to 1-2Mbps after a while even for models as small as 2GB.

Is there a way to prevent this?

[FL33TW00D]

Do you have a minimal reproducible example? I've just tested downloading a model on my end the speed looks good!

[CodeWithShreyans]

I'm using it in my app https://shreyans.sh/othree and I've seen the problem myself and gotten feedback about it from other users too.

[CodeWithShreyans]

I could absolutely be doing something wrong here as I'm pretty new to iOS development but all the download code is from fullmoon: https://github.com/mainframecomputer/fullmoon-ios/blob/main/fullmoon/Models/LLMEvaluator.swift#L68

[CodeWithShreyans]

Does the downloader not make multiple server connections?
Does it support something like HF_HUB_ENABLE_HF_TRANSFER=1 that the cli does?

[pcuenca]

No, not yet. But I see no reason why download speed should decrease after it has started. Could it be a temporary networking issue? Of course there could also be a bug somewhere.

[CodeWithShreyans]

I tried a workaround where I created a custom HubApi instance with a custom endpoint I created that routes the requests to model files hosted on my own S3 bucket. It improved it somewhat but the issue still persisted. I tried both the default and workaround on multiple fast networks with the same results. This tells me it's likely a bug in the package.

[pcuenca]

Interesting. cc @ardaatahan, just for info.

[ZachNagengast]

@CodeWithShreyans Can you share any code you're using to call this, or even the URL you are downloading from? It's possible that there is some missing metadata that isn't being handled well. Perhaps if you can put a breakpoint here

swift-transformers/Sources/Hub/HubApi.swift

Lines 495 to 497 in 1d26ce0

	func getFileMetadata(url: URL) async throws -> FileMetadata {
	let (_, response) = try await httpHead(for: url)
	let location = response.statusCode == 302 ? response.value(forHTTPHeaderField: "Location") : response.url?.absoluteString

It will tell you the metadata that's coming back and we can see whats missing. You may also want to try using main branch and see if anything is improved.

[CodeWithShreyans]

[ZachNagengast]

@CodeWithShreyans can you provide any more info about the metadata specifically? Also if possible anymore detailed info about the slower speeds you are seeing compared to other similar downloaders like the python version. Would like to isolate whether you're experiencing speeds that are out of parity with typical hf download speeds, S3 won’t be an exactly apples to apples comparison.

[ardaatahan]

Hey @CodeWithShreyans, just following up - can you share more details on the metadata and slower speeds? We're trying to reproduce the issue so we can move forward with merging the offline mode PR. Appreciate any extra context!

cc: @FL33TW00D

[CodeWithShreyans]

The bigger problem doesn't seem to be the speed but the instability.

On the same fast network, sometimes it downloads fine, sometimes it takes a while, sometimes it just gets stuck (like it did just now).

Tried out the python downloader on the same network, took about 1 min each try and never got stuck.

[CodeWithShreyans]

let modelContainer = try await LLMModelFactory.shared.loadContainer(
    configuration: model
) {
    [modelConfiguration] progress in
        Task { @MainActor in
            self.modelInfo =
              "Downloading \(modelConfiguration.name): \(Int(progress.fractionCompleted * 100))%"
            self.progress = progress.fractionCompleted
    }
}

I call loadContainer() from mlx-libraries which then calls this:

public func _loadContainer(
    hub: HubApi = HubApi(), configuration: ModelConfiguration,
    progressHandler: @Sendable @escaping (Progress) -> Void = { _ in }
) async throws -> ModelContainer {
    let context = try await _load(
        hub: hub, configuration: configuration, progressHandler: progressHandler)
    return ModelContainer(context: context)
}


    public func _load(
        hub: HubApi, configuration: ModelConfiguration,
        progressHandler: @Sendable @escaping (Progress) -> Void
    ) async throws -> ModelContext {
        // download weights and config
        let modelDirectory = try await downloadModel(
            hub: hub, configuration: configuration, progressHandler: progressHandler)

        // load the generic config to unerstand which model and how to load the weights
        let configurationURL = modelDirectory.appending(component: "config.json")
        let baseConfig = try JSONDecoder().decode(
            BaseConfiguration.self, from: Data(contentsOf: configurationURL))
        let model = try typeRegistry.createModel(
            configuration: configurationURL, modelType: baseConfig.modelType)

        // apply the weights to the bare model
        try loadWeights(
            modelDirectory: modelDirectory, model: model, quantization: baseConfig.quantization)

        let tokenizer = try await loadTokenizer(configuration: configuration, hub: hub)

        return .init(
            configuration: configuration, model: model,
            processor: LLMUserInputProcessor(tokenizer: tokenizer, configuration: configuration),
            tokenizer: tokenizer)
    }


public func downloadModel(
    hub: HubApi, configuration: ModelConfiguration,
    progressHandler: @Sendable @escaping (Progress) -> Void
) async throws -> URL {
    do {
        switch configuration.id {
        case .id(let id):
            // download the model weights
            let repo = Hub.Repo(id: id)
            let modelFiles = ["*.safetensors", "*.json"]
            return try await hub.snapshot(
                from: repo, matching: modelFiles, progressHandler: progressHandler)

        case .directory(let directory):
            return directory
        }

    } catch Hub.HubClientError.authorizationRequired {
        // an authorizationRequired means (typically) that the named repo doesn't exist on
        // on the server so retry with local only configuration
        return configuration.modelDirectory(hub: hub)

    } catch {
        let nserror = error as NSError
        if nserror.domain == NSURLErrorDomain && nserror.code == NSURLErrorNotConnectedToInternet {
            // Error Domain=NSURLErrorDomain Code=-1009 "The Internet connection appears to be offline."
            // fall back to the local directory
            return configuration.modelDirectory(hub: hub)
        } else {
            throw error
        }
    }