SPB Git

spb/forge-studio Public

The Instruments of LLM training — a native macOS cockpit for Forge. Train language models from scratch on Apple Silicon without a terminal.

Swift 95.7% Shell 4.3%
15.3 KB · 335 lines swift
Raw Blame History
1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// New Run editor: preset picker, the high-traffic model/train fields, the4// live derived panel (params, tokens/step, epochs, memory badge) and the LR5// schedule preview. Inline validation gates the Start button. (The long6// tail of variant knobs lives in the "Avancé" JSON editor — full form in M2.)7import Charts8import SwiftUI910struct NewRunSheet: View {11    @Environment(AppModel.self) private var app12    @Environment(\.dismiss) private var dismiss13    @State private var config = ForgeConfig()14    @State private var runName = "run"15    @State private var datasetPath = ""16    @State private var startError: String?1718    private var datasets: [Dataset] { app.datasets() }19    private var dataset: Dataset? { datasets.first { $0.path == datasetPath } }20    private var errors: [String] {21        var e = config.validationErrors22        if datasetPath.isEmpty { e.append("choisir un dataset") }23        if let v = dataset?.vocabSize, v != config.model.vocabSize {24            e.append("vocab_size (\(config.model.vocabSize)) ≠ tokenizer du dataset (\(v))")25        }26        return e27    }2829    var body: some View {30        VStack(spacing: 0) {31            HStack {32                Text("Nouveau run").font(.title2.weight(.semibold))33                Spacer()34                presetMenu35            }36            .padding()37            Divider()38            HSplitView {39                Form {40                    Section("Run") {41                        TextField("Nom", text: $runName)42                        Picker("Dataset", selection: $datasetPath) {43                            Text("—").tag("")44                            ForEach(datasets) { ds in45                                Text("\(ds.name) · vocab \(ds.vocabSize ?? 0) · \((ds.trainTokens ?? 0).formatted(.number.notation(.compactName))) tokens")46                                    .tag(ds.path)47                            }48                        }49                        .onChange(of: datasetPath) {50                            if let v = dataset?.vocabSize { config.model.vocabSize = v }51                        }52                    }53                    Section("Modèle") {54                        intField("n_layers", $config.model.nLayers)55                        intField("d_model", $config.model.dModel)56                        intField("n_heads", $config.model.nHeads)57                        intField("n_kv_heads (GQA)", $config.model.nKvHeads)58                        intField("d_ff", $config.model.dFf)59                        intField("context_length", $config.model.contextLength)60                        Picker("activation", selection: $config.model.activation) {61                            ForEach(["swiglu", "gelu", "relu2"], id: \.self) { Text($0) }62                        }63                        Picker("norm", selection: $config.model.norm) {64                            ForEach(["rmsnorm", "layernorm"], id: \.self) { Text($0) }65                        }66                        Toggle("tied_embeddings", isOn: $config.model.tiedEmbeddings)67                        Toggle("qk_norm", isOn: $config.model.qkNorm)68                        intField("sliding_window (0 = full)", $config.model.slidingWindow)69                    }70                    Section("Attention & positions") {71                        Toggle("use_rope", isOn: $config.model.useRope)72                        doubleField("rope_theta", $config.model.ropeTheta)73                        doubleField("rope_theta_global (0 = idem)", $config.model.ropeThetaGlobal)74                        intField("sliding_global_every (Gemma3: 6)", $config.model.slidingGlobalEvery)75                        intField("nope_every (SmolLM3: 4)", $config.model.nopeEvery)76                        intField("head_dim (0 = auto)", $config.model.headDimOverride)77                        Toggle("attention_bias (Qwen2.5)", isOn: $config.model.attentionBias)78                        doubleField("attn_softcap (Gemma2: 50)", $config.model.attnSoftcap)79                        doubleField("rope_scale_factor (llama3: 32, 0 = off)",80                                    $config.model.ropeScaleFactor)81                        if config.model.ropeScaleFactor > 0 {82                            doubleField("rope_scale_low", $config.model.ropeScaleLow)83                            doubleField("rope_scale_high", $config.model.ropeScaleHigh)84                            intField("rope_scale_orig_ctx", $config.model.ropeScaleOrigCtx)85                        }86                    }87                    Section("Normalisation & sorties") {88                        Picker("norm_placement", selection: $config.model.normPlacement) {89                            ForEach(["pre", "post", "sandwich"], id: \.self) { Text($0) }90                        }91                        doubleField("norm_eps", $config.model.normEps)92                        doubleField("final_softcap (Gemma2: 30)", $config.model.finalSoftcap)93                        Toggle("scale_embeddings (Gemma)", isOn: $config.model.scaleEmbeddings)94                        Picker("quant (QAT)", selection: $config.model.quant) {95                            ForEach(["none", "int8", "ternary"], id: \.self) { Text($0) }96                        }97                    }98                    Section("Mixture of Experts") {99                        intField("n_experts (0 = dense)", $config.model.nExperts)100                        if config.model.nExperts > 0 {101                            intField("moe_top_k", $config.model.moeTopK)102                            intField("n_shared_experts (V3: 1)", $config.model.nSharedExperts)103                            Picker("moe_scoring", selection: $config.model.moeScoring) {104                                ForEach(["softmax", "sigmoid"], id: \.self) { Text($0) }105                            }106                            Toggle("moe_norm_topk", isOn: $config.model.moeNormTopk)107                            doubleField("routed_scaling_factor (V3: 2.5)",108                                        $config.model.routedScalingFactor)109                            intField("moe_d_ff (0 = d_ff)", $config.model.moeDFf)110                            intField("first_k_dense (V3: 3)", $config.model.firstKDense)111                            doubleField("moe_bias_gamma (noaux, V3: 0.001)",112                                        $config.model.moeBiasGamma)113                            doubleField("moe_aux_weight", $config.model.moeAuxWeight)114                        }115                    }116                    Section("Entraînement") {117                        doubleField("lr", $config.train.lr)118                        intField("max_steps", $config.train.maxSteps)119                        intField("warmup_steps", $config.train.warmupSteps)120                        Picker("schedule", selection: $config.train.schedule) {121                            ForEach(["cosine", "wsd"], id: \.self) { Text($0) }122                        }123                        Picker("optimizer", selection: $config.train.optimizer) {124                            ForEach(["adamw", "muon"], id: \.self) { Text($0) }125                        }126                        intField("batch_size (micro)", $config.train.batchSize)127                        intField("grad_accum_steps", $config.train.gradAccumSteps)128                        intField("checkpoint_every", $config.train.checkpointEvery)129                        intField("eval_every", $config.train.evalEvery)130                    }131                    Section("Optimiseur avancé") {132                        if config.train.optimizer == "muon" {133                            doubleField("muon_lr", $config.train.muonLr)134                            doubleField("muon_momentum", $config.train.muonMomentum)135                        }136                        if config.train.schedule == "wsd" {137                            doubleField("wsd_decay_frac", $config.train.wsdDecayFrac)138                        }139                        doubleField("min_lr_ratio", $config.train.minLrRatio)140                        doubleField("beta1", $config.train.beta1)141                        doubleField("beta2", $config.train.beta2)142                        doubleField("weight_decay", $config.train.weightDecay)143                        doubleField("grad_clip (0 = off)", $config.train.gradClip)144                        intField("eval_batches", $config.train.evalBatches)145                        intField("seed", $config.train.seed)146                        Toggle("forge_save (.forge à chaque checkpoint)",147                               isOn: $config.train.forgeSave)148                        Picker("forge_dtype", selection: $config.train.forgeDtype) {149                            ForEach(["f32", "f16", "bf16"], id: \.self) { Text($0) }150                        }151                    }152                }153                .formStyle(.grouped)154                .frame(minWidth: 380)155156                DerivedPanel(config: config, dataset: dataset)157                    .frame(minWidth: 300)158                    .padding()159            }160            Divider()161            HStack {162                if let first = errors.first {163                    Label(first, systemImage: "exclamationmark.triangle")164                        .foregroundStyle(.red)165                        .font(.callout)166                }167                if let startError {168                    Text(startError).foregroundStyle(.red).font(.callout)169                }170                Spacer()171                Button("Annuler") { dismiss() }172                Button("Démarrer l'entraînement") {173                    Task {174                        do {175                            config.model.name = runName176                            try await app.supervisor.start(177                                config: config, datasetPath: datasetPath,178                                name: runName)179                            dismiss()180                        } catch {181                            startError = error.localizedDescription182                        }183                    }184                }185                .keyboardShortcut(.defaultAction)186                .disabled(!errors.isEmpty)187            }188            .padding()189        }190        .frame(minWidth: 860, minHeight: 620)191        .onAppear {192            if let first = datasets.first { datasetPath = first.path }193        }194    }195196    private var presetMenu: some View {197        HStack {198            Menu("Presets") {199                ForEach(Presets.all, id: \.name) { preset in200                    Button(preset.name) {201                        config = preset.config202                        runName = preset.config.model.name203                        if let v = dataset?.vocabSize { config.model.vocabSize = v }204                    }205                }206            }207            .frame(width: 110)208            Button("Importer…") { importConfig() }209            Button("Exporter…") { exportConfig() }210        }211    }212213    private func importConfig() {214        let panel = NSOpenPanel()215        panel.allowedContentTypes = [.json]216        if panel.runModal() == .OK, let url = panel.url,217           let cfg = try? ForgeConfig.load(from: url) {218            config = cfg219            runName = cfg.model.name220        }221    }222223    private func exportConfig() {224        let panel = NSSavePanel()225        panel.allowedContentTypes = [.json]226        panel.nameFieldStringValue = "\(runName).json"227        if panel.runModal() == .OK, let url = panel.url {228            var cfg = config229            cfg.model.name = runName230            try? cfg.exportJSON().write(to: url)231        }232    }233234    private func intField(_ label: String, _ value: Binding<Int>) -> some View {235        TextField(label, value: value, format: .number)236    }237238    private func doubleField(_ label: String, _ value: Binding<Double>) -> some View {239        TextField(label, value: value, format: .number.precision(.significantDigits(1...6)))240    }241}242243struct DerivedPanel: View {244    let config: ForgeConfig245    let dataset: Dataset?246247    var body: some View {248        VStack(alignment: .leading, spacing: 14) {249            Text("Dérivés").font(.headline)250            derived("Paramètres",251                    config.model.paramCount.formatted(.number.notation(.compactName)))252            derived("Tokens/step", config.tokensPerStep.formatted())253            derived("Tokens totaux",254                    config.totalTokens.formatted(.number.notation(.compactName)))255            if let t = dataset?.trainTokens, t > 0 {256                derived("Epochs sur le dataset",257                        String(format: "%.2f", Double(config.totalTokens) / Double(t)))258            }259            memoryBadge260            Divider()261            Text("Schedule LR").font(.headline)262            Chart {263                let steps = stride(from: 0, to: config.train.maxSteps,264                                   by: max(1, config.train.maxSteps / 200))265                ForEach(Array(steps), id: \.self) { s in266                    LineMark(x: .value("step", s),267                             y: .value("lr", config.train.lrAt(step: s)))268                        .foregroundStyle(.purple)269                }270            }271            .frame(height: 140)272            Spacer()273        }274    }275276    private var memoryBadge: some View {277        let bytes = SystemInfo.estimatedTrainingBytes(config: config)278        let machine = SystemInfo.memoryBytes279        let gb = Double(bytes) / 1e9280        let ok = bytes < machine * 8 / 10281        return HStack {282            Image(systemName: ok ? "memorychip" : "exclamationmark.triangle.fill")283            Text(String(format: "~%.1f GB estimés / %.0f GB unifiés", gb,284                        Double(machine) / 1e9))285        }286        .font(.callout)287        .foregroundStyle(ok ? Color.secondary : Color.orange)288    }289290    private func derived(_ label: String, _ value: String) -> some View {291        HStack {292            Text(label).foregroundStyle(.secondary)293            Spacer()294            Text(value).monospacedDigit()295        }296        .font(.callout)297    }298}299300enum Presets {301    struct Preset {302        let name: String303        let config: ForgeConfig304    }305306    // Mirrors configs/ in the forge repo (gpt-10m-1epoch, gpt-50m…).307    static let all: [Preset] = {308        var p10 = ForgeConfig()309        p10.model.name = "gpt-10m"310        p10.model.nLayers = 6; p10.model.dModel = 384; p10.model.nHeads = 6311        p10.model.nKvHeads = 6; p10.model.dFf = 1024; p10.model.vocabSize = 4096312        p10.model.contextLength = 512313        p10.train.lr = 6e-4; p10.train.warmupSteps = 58; p10.train.maxSteps = 584314        p10.train.batchSize = 64; p10.train.gradAccumSteps = 1315        p10.train.checkpointEvery = 200; p10.train.evalEvery = 100316317        var p50 = ForgeConfig()318        p50.model.name = "gpt-50m"319        p50.model.nLayers = 10; p50.model.dModel = 640; p50.model.nHeads = 10320        p50.model.nKvHeads = 10; p50.model.dFf = 1728; p50.model.vocabSize = 4096321        p50.model.contextLength = 1024322        p50.train.lr = 5e-4; p50.train.warmupSteps = 117; p50.train.maxSteps = 1170323        p50.train.batchSize = 8; p50.train.gradAccumSteps = 8324        p50.train.checkpointEvery = 200; p50.train.evalEvery = 100325326        var muon = p50327        muon.model.name = "gpt-50m-muon"328        muon.train.optimizer = "muon"; muon.train.schedule = "wsd"329330        return [Preset(name: "gpt-10m (1 epoch)", config: p10),331                Preset(name: "gpt-50m", config: p50),332                Preset(name: "gpt-50m Muon+WSD", config: muon)]333    }()334}335