Implement GCS Checkpoint

architectures/centralized/client/src/app.rs

@@ -2,6 +2,8 @@ use anyhow::{Error, Result};
 use bytemuck::Zeroable;
 use hf_hub::Repo;
 use psyche_centralized_shared::{ClientId, ClientToServerMessage, ServerToClientMessage};
+use psyche_client::HubUploadInfo;
+use psyche_client::UploadInfo;
 use psyche_client::{
     Client, ClientTUI, ClientTUIState, NC, RunInitConfig, TrainArgs, read_identity_secret_key,
 };
@@ -29,7 +31,7 @@ pub type TabsData = <Tabs as CustomWidget>::Data;
 pub enum ToSend {
     Witness(Box<OpportunisticData>),
     HealthCheck(HealthChecks<ClientId>),
-    Checkpoint(model::HubRepo),
+    Checkpoint(model::Checkpoint),
 }
 
 struct Backend {
@@ -67,7 +69,7 @@ impl WatcherBackend<ClientId> for Backend {
         Ok(())
     }
 
-    async fn send_checkpoint(&mut self, checkpoint: model::HubRepo) -> Result<()> {
+    async fn send_checkpoint(&mut self, checkpoint: model::Checkpoint) -> Result<()> {
         self.tx.send(ToSend::Checkpoint(checkpoint))?;
         Ok(())
     }
@@ -176,18 +178,19 @@ impl App {
     ) -> Result<()> {
         // sanity checks
         if let Some(checkpoint_config) = &state_options.checkpoint_config {
-            if let Some(hub_upload) = &checkpoint_config.hub_upload {
+            if let Some(UploadInfo::Hub(HubUploadInfo {
+                hub_repo,
+                hub_token,
+            })) = &checkpoint_config.upload_info
+            {
                 let api = hf_hub::api::tokio::ApiBuilder::new()
-                    .with_token(Some(hub_upload.hub_token.clone()))
+                    .with_token(Some(hub_token.clone()))
                     .build()?;
-                let repo_api = api.repo(Repo::new(
-                    hub_upload.hub_repo.clone(),
-                    hf_hub::RepoType::Model,
-                ));
+                let repo_api = api.repo(Repo::new(hub_repo.clone(), hf_hub::RepoType::Model));
                 if !repo_api.is_writable().await {
                     anyhow::bail!(
                         "Checkpoint upload repo {} is not writable with the passed API key.",
-                        hub_upload.hub_repo
+                        hub_repo
                     )
                 }
             }

architectures/centralized/server/src/app.rs

@@ -9,7 +9,8 @@ use psyche_coordinator::{
 
 use psyche_core::{FixedVec, Shuffle, SizedIterator, TokenSize};
 use psyche_data_provider::{
-    DataProviderTcpServer, DataServerTui, LocalDataProvider, download_model_repo_async,
+    DataProviderTcpServer, DataServerTui, LocalDataProvider, download_model_from_gcs_async,
+    download_model_repo_async,
 };
 use psyche_network::{ClientNotification, TcpServer};
 use psyche_tui::{
@@ -80,7 +81,7 @@ impl psyche_watcher::Backend<ClientId> for ChannelCoordinatorBackend {
         bail!("Server does not send health checks");
     }
 
-    async fn send_checkpoint(&mut self, _checkpoint: model::HubRepo) -> Result<()> {
+    async fn send_checkpoint(&mut self, _checkpoint: model::Checkpoint) -> Result<()> {
         bail!("Server does not send checkpoints");
     }
 }
@@ -201,9 +202,15 @@ impl App {
                             Checkpoint::Dummy(_) => {
                                 // ok!
                             }
-                            Checkpoint::P2P(_) => {
+                            Checkpoint::P2P(_) | Checkpoint::P2PGcs(_) => {
                                 bail!("Can't start up a run with a P2P checkpoint.")
                             }
+                            Checkpoint::Gcs(gcs_repo) => {
+                                let bucket: String = (&gcs_repo.bucket).into();
+                                let prefix: Option<String> =
+                                    gcs_repo.prefix.map(|p| (&p).into());
+                                download_model_from_gcs_async(&bucket, prefix.as_deref()).await?;
+                            }
                         }
 
                         let server_addr: SocketAddr = String::from(url).parse().map_err(|e| {

architectures/centralized/shared/src/protocol.rs

@@ -15,7 +15,7 @@ pub enum ClientToServerMessage {
     Join { run_id: String },
     Witness(Box<OpportunisticData>),
     HealthCheck(HealthChecks<ClientId>),
-    Checkpoint(model::HubRepo),
+    Checkpoint(model::Checkpoint),
 }
 
 #[derive(Serialize, Deserialize, Debug, Clone)]

architectures/decentralized/solana-client/src/main.rs

@@ -283,34 +283,50 @@ async fn async_main() -> Result<()> {
                     bail!("Model is not an LLM, unsure how to predownload.");
                 };
 
-                let checkpoint = match model_config.checkpoint {
+                match model_config.checkpoint {
                     Checkpoint::Ephemeral => {
                         bail!("Can't predownload model with ephemeral checkpoint.")
                     }
                     Checkpoint::Dummy(hub_repo)
                     | Checkpoint::Hub(hub_repo)
-                    | Checkpoint::P2P(hub_repo) => hub_repo,
+                    | Checkpoint::P2P(hub_repo) => {
+                        let repo_id = hub_repo.repo_id.to_string();
+                        let revision = hub_repo.revision.map(|s| s.to_string());
+                        println!(
+                            "Predownloading model {repo_id} revision {}",
+                            revision.as_ref().unwrap_or(&"main".to_string())
+                        );
+
+                        let hub_read_token = std::env::var("HF_TOKEN").ok();
+                        let cache_folder = None; // Uses HF_HOME env var
+
+                        psyche_data_provider::download_model_repo_async(
+                            &repo_id,
+                            revision,
+                            cache_folder,
+                            hub_read_token,
+                            Some(hub_max_concurrent_downloads),
+                            true,
+                        )
+                        .await?;
+                    }
+                    Checkpoint::Gcs(gcs_repo) | Checkpoint::P2PGcs(gcs_repo) => {
+                        let bucket = gcs_repo.bucket.to_string();
+                        let prefix: Option<String> = gcs_repo.prefix.map(|p| p.to_string());
+                        println!(
+                            "Predownloading model from gs://{}/{}",
+                            bucket,
+                            prefix.as_deref().unwrap_or("")
+                        );
+
+                        psyche_data_provider::download_model_from_gcs_async(
+                            &bucket,
+                            prefix.as_deref(),
+                        )
+                        .await?;
+                    }
                 };
 
-                let repo_id = checkpoint.repo_id.to_string();
-                let revision = checkpoint.revision.map(|s| s.to_string());
-                println!(
-                    "Predownloading model {repo_id} revision {}",
-                    revision.as_ref().unwrap_or(&"main".to_string())
-                );
-
-                let hub_read_token = std::env::var("HF_TOKEN").ok();
-                let cache_folder = None; // Uses HF_HOME env var
-
-                psyche_data_provider::download_model_repo_async(
-                    &repo_id,
-                    revision,
-                    cache_folder,
-                    hub_read_token,
-                    Some(hub_max_concurrent_downloads),
-                    true,
-                )
-                .await?;
                 println!("Model predownloaded successfully.");
             }
 

architectures/decentralized/solana-common/src/backend.rs

@@ -19,7 +19,8 @@ use anchor_client::{
 };
 use anyhow::{Context, Result, anyhow};
 use futures_util::StreamExt;
-use psyche_coordinator::{CommitteeProof, Coordinator, HealthChecks, model::HubRepo};
+use psyche_coordinator::model::{self, Checkpoint};
+use psyche_coordinator::{CommitteeProof, Coordinator, HealthChecks};
 use psyche_core::IntegrationTestLogMarker;
 use psyche_watcher::{Backend as WatcherBackend, OpportunisticData};
 use solana_account_decoder_client_types::{UiAccount, UiAccountEncoding};
@@ -334,7 +335,7 @@ impl SolanaBackend {
         &self,
         coordinator_instance: Pubkey,
         coordinator_account: Pubkey,
-        repo: HubRepo,
+        repo: Checkpoint,
     ) {
         let user = self.get_payer();
         let instruction = instructions::coordinator_checkpoint(
@@ -604,7 +605,7 @@ impl WatcherBackend<psyche_solana_coordinator::ClientId> for SolanaBackendRunner
         Ok(())
     }
 
-    async fn send_checkpoint(&mut self, checkpoint: HubRepo) -> Result<()> {
+    async fn send_checkpoint(&mut self, checkpoint: model::Checkpoint) -> Result<()> {
         self.backend
             .send_checkpoint(self.instance, self.account, checkpoint);
         Ok(())

architectures/decentralized/solana-common/src/instructions.rs

@@ -206,7 +206,7 @@ pub fn coordinator_checkpoint(
     coordinator_instance: &Pubkey,
     coordinator_account: &Pubkey,
     user: &Pubkey,
-    repo: psyche_coordinator::model::HubRepo,
+    repo: psyche_coordinator::model::Checkpoint,
 ) -> Instruction {
     anchor_instruction(
         psyche_solana_coordinator::ID,

architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/instance_state.rs

@@ -10,7 +10,7 @@ use psyche_coordinator::RunState;
 use psyche_coordinator::SOLANA_MAX_STRING_LEN;
 use psyche_coordinator::TickResult;
 use psyche_coordinator::Witness;
-use psyche_coordinator::model::HubRepo;
+use psyche_coordinator::model::Checkpoint;
 use psyche_coordinator::model::Model;
 use psyche_core::FixedString;
 use psyche_core::SmallBoolean;
@@ -389,7 +389,11 @@ impl CoordinatorInstanceState {
         self.tick()
     }
 
-    pub fn checkpoint(&mut self, payer: &Pubkey, repo: HubRepo) -> Result<()> {
+    pub fn checkpoint(
+        &mut self,
+        payer: &Pubkey,
+        repo: Checkpoint,
+    ) -> Result<()> {
         // O(n) on clients, reconsider
         let id = self.clients_state.find_signer(payer)?;
         let index = self

architectures/decentralized/solana-coordinator/programs/solana-coordinator/src/lib.rs

@@ -21,7 +21,7 @@ use psyche_coordinator::Witness;
 use psyche_coordinator::WitnessBloom;
 use psyche_coordinator::WitnessMetadata;
 use psyche_coordinator::WitnessProof;
-use psyche_coordinator::model::{HubRepo, Model};
+use psyche_coordinator::model::Model;
 use psyche_core::MerkleRoot;
 use serde::Deserialize;
 use serde::Serialize;
@@ -313,7 +313,7 @@ pub mod psyche_solana_coordinator {
 
     pub fn checkpoint(
         ctx: Context<PermissionlessCoordinatorAccounts>,
-        repo: HubRepo,
+        repo: psyche_coordinator::model::Checkpoint,
     ) -> Result<()> {
         let mut account = ctx.accounts.coordinator_account.load_mut()?;
         account.increment_nonce();

config/solana-test/light-config-gcs.toml

@@ -0,0 +1,44 @@
+[config]
+warmup_time = 30
+cooldown_time = 30
+epoch_time = 60
+max_round_train_time = 15
+round_witness_time = 1
+min_clients = 1
+init_min_clients = 1
+verification_percent = 0
+witness_nodes = 0
+global_batch_size_start = 8
+global_batch_size_end = 8
+global_batch_size_warmup_tokens = 0
+total_steps = 25000
+waiting_for_members_extra_time = 3
+
+[model.LLM]
+architecture = "HfLlama"
+data_type = "Pretraining"
+max_seq_len = 2048
+cold_start_warmup_steps = 0
+
+[model.LLM.checkpoint.Gcs]
+bucket = "llama220minit"
+
+[model.LLM.data_location.Http]
+token_size_in_bytes = "TwoBytes"
+shuffle = "DontShuffle"
+[model.LLM.data_location.Http.location.Gcp]
+bucket_name = "nous-pretraining-public-us"
+filter_directory = "fineweb-edu-tokenized-llama2"
+
+[model.LLM.lr_schedule.Cosine]
+base_lr = 4.0e-4
+warmup_steps = 250
+warmup_init_lr = 0.0
+total_steps = 25000
+final_lr = 4.0e-5
+[model.LLM.optimizer.Distro]
+clip_grad_norm = 1.0
+compression_decay = 0.999
+compression_chunk = 64
+compression_topk = 8
+quantize_1bit = true

nix/lib.nix

@@ -38,6 +38,7 @@ let
       python312
       pkg-config
       perl
+      cargo-nextest
     ];
 
     buildInputs =

psyche-book/src/SUMMARY.md

@@ -14,6 +14,7 @@
   - [General workflow](./explain/general-workflow.md)
   - [Data provider](./explain/data-provider.md)
   - [Model sharing](./explain/model-sharing.md)
+  - [GCS Checkpoints](./explain/gcs-checkpoints.md)
   - [Rewards](./explain/rewards.md)
   - [Glossary](./explain/glossary.md)