Concurrent Checkpointing

Concurrent Checkpointing is a feature that asynchronously saves checkpoints even during Forward/Backward operations, unlike the traditional method. Therefore, by using this feature you can reduce the overhead of checkpoint saving, efficiently shorten the overall training time, and automatically save the training progress that might be lost if an unexpected interruption occurs.

Concurrent Checkpointing Overview

The Training Jobs provided by Simple AI Training include On-Demand Training type and Spot Training type. Concurrent Checkpointing can be used in both types, but the scope of its functionality differs by type. The scope of use for each type is as follows.

Using Concurrent Checkpointing

Preliminary preparation: Write script

The user can use the save method of Trainer and Concurrent Checkpoint simultaneously.

Spot Training Usage

The script example when using Spot Training is as follows.

|language = python | title = Training Script Example | collapse = true

// Written based on transformers==5.10.2.

import os
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling
)
from datasets import load_dataset
import json
from datastates.llm import DecoratedCheckpointing
import argparse
import logging
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        --local_rank
        type=int,
        default=-1,
        help="local rank passed from distributed launcher (Deepspeed, torchrun, etc.)"
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()

    model_path="/root/.cache/huggingface/hub/models--meta-llama--Llama-3.2-1B/snapshots/4e20de362430cd3b72f300e6b0f18e50e7166e08"

    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)

    # Set pad token to EOS if not already defined
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Load WikiText-2 dataset
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1",cache_dir="/root/.cache/huggingface/datasets")

    # Tokenization function
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=128,
            padding="max_length"
        )

    # Tokenize the dataset
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=["text"]
    )

    train_dataset = tokenized_dataset["train"]
    valid_dataset = tokenized_dataset["validation"]

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False  # Causal LM (not masked LM)
    )

    script_directory = os.path.dirname(os.path.abspath(__file__))
    ds_config_path = os.path.join(script_directory, "ds_config.json")

    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=3,
        per_device_train_batch_size=2,  # Adjust based on GPU memory
        gradient_accumulation_steps=4,  # Effective batch size = batch_size * gradient_accumulation_steps
        save_strategy="steps",
        save_steps=200,
        logging_steps=2,
        eval_strategy="steps",
        eval_steps=100,
        bf16=True,  # Enable BF16 mixed precision (use fp16 if unsupported)
        deepspeed=ds_config_path,  # Path to DeepSpeed config file
        report_to="none",
    )

    model = AutoModelForCausalLM.from_pretrained( model_path,  local_files_only=True, low_cpu_mem_usage=True, device_map=None)

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,  # Optional: validation set for evaluation
        processing_class=tokenizer,
        data_collator=data_collator,
    )

       # ADD configuration for Concurrent CHECKPOINT ENGINE
    config = {
        "host_cache_size": 50,
        "parser_threads": 1,
        "pin_host_cache": True,
        "trainer": trainer,
    }

    ckpt_engine = DecoratedCheckpointing(runtime_config=config, rank=args.local_rank)

resume_from_checkpoint=False

if os.getenv("CKPT_LAST_STEP") != None :
resume_from_checkpoint=True

trainer.train(resume_from_checkpoint=resume_from_checkpoint)

// Written based on transformers==5.10.2.

import os
import torch
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling
)
from datasets import load_dataset
import json
from datastates.llm import DecoratedCheckpointing
import argparse
import logging
import time

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        --local_rank
        type=int,
        default=-1,
        help="local rank passed from distributed launcher (Deepspeed, torchrun, etc.)"
    )
    return parser.parse_args()

if __name__ == "__main__":
    args = parse_args()

    model_path="/root/.cache/huggingface/hub/models--meta-llama--Llama-3.2-1B/snapshots/4e20de362430cd3b72f300e6b0f18e50e7166e08"

    # Load tokenizer and model
    tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True)

    # Set pad token to EOS if not already defined
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # Load WikiText-2 dataset
    dataset = load_dataset("wikitext", "wikitext-2-raw-v1",cache_dir="/root/.cache/huggingface/datasets")

    # Tokenization function
    def tokenize_function(examples):
        return tokenizer(
            examples["text"],
            truncation=True,
            max_length=128,
            padding="max_length"
        )

    # Tokenize the dataset
    tokenized_dataset = dataset.map(
        tokenize_function,
        batched=True,
        remove_columns=["text"]
    )

    train_dataset = tokenized_dataset["train"]
    valid_dataset = tokenized_dataset["validation"]

    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False  # Causal LM (not masked LM)
    )

    script_directory = os.path.dirname(os.path.abspath(__file__))
    ds_config_path = os.path.join(script_directory, "ds_config.json")

    training_args = TrainingArguments(
        output_dir="./results",
        num_train_epochs=3,
        per_device_train_batch_size=2,  # Adjust based on GPU memory
        gradient_accumulation_steps=4,  # Effective batch size = batch_size * gradient_accumulation_steps
        save_strategy="steps",
        save_steps=200,
        logging_steps=2,
        eval_strategy="steps",
        eval_steps=100,
        bf16=True,  # Enable BF16 mixed precision (use fp16 if unsupported)
        deepspeed=ds_config_path,  # Path to DeepSpeed config file
        report_to="none",
    )

    model = AutoModelForCausalLM.from_pretrained( model_path,  local_files_only=True, low_cpu_mem_usage=True, device_map=None)

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=valid_dataset,  # Optional: validation set for evaluation
        processing_class=tokenizer,
        data_collator=data_collator,
    )

       # ADD configuration for Concurrent CHECKPOINT ENGINE
    config = {
        "host_cache_size": 50,
        "parser_threads": 1,
        "pin_host_cache": True,
        "trainer": trainer,
    }

    ckpt_engine = DecoratedCheckpointing(runtime_config=config, rank=args.local_rank)

resume_from_checkpoint=False

if os.getenv("CKPT_LAST_STEP") != None :
resume_from_checkpoint=True

trainer.train(resume_from_checkpoint=resume_from_checkpoint)

다음 절차의 예시를 참고하여 스크립트를 작성하세요.

1. Import Concurrent CHECKPOINT

from datastates.llm import DecoratedCheckpointing
...

2. ADD configuration for Concurrent CHECKPOINT ENGINE

config = {
       "host_cache_size": 50,
       "parser_threads": 1,
       "pin_host_cache": True,
       "trainer": trainer,
}

3. Initialize Concurrent CHECKPOINT ENGINE

ckpt_engine = DecoratedCheckpointing(runtime_config=config, rank=args.local_rank)

4. Set Concurrent CHECKPOINT ENGINE parameter

resume_from_checkpoint=False

if os.getenv("CKPT_LAST_STEP") != None :
       resume_from_checkpoint=True

trainer.train(resume_from_checkpoint=resume_from_checkpoint)

Concurrent Checkpoint save path

The save path is based by default on the TrainingArgument’s output_dir. It is stored under the concurrent_checkpoint directory in the subpath of output_dir.

Using On-Demand Training

The method for using On-Demand Training is similar to that of Spot Training.

Initial training

Write the script by referring to the example of the following procedure.

1. Import Concurrent CHECKPOINT

from datastates.llm import DecoratedCheckpointing
...

2. ADD configuration for Concurrent CHECKPOINT ENGINE

config = {
       "host_cache_size": 50,
       "parser_threads": 1,
       "pin_host_cache": True,
       "trainer": trainer,
}

3. Initialize Concurrent CHECKPOINT ENGINE

ckpt_engine = DecoratedCheckpointing(runtime_config=config, rank=args.local_rank)

4. Set Concurrent CHECKPOINT ENGINE parameter

trainer.train(resume_from_checkpoint=False)

When manually activated

If a valid checkpoint is found in the output_dir path specified during the initial training, the checkpoint path is directly specified during model initialization when re-running the training. Or, when re-running training, specify the same output_dir as before and configure as follows to automatically load the latest checkpoint.

Write the script by referring to the example of the following procedure.

1. Import Concurrent CHECKPOINT

from datastates.llm import DecoratedCheckpointing
...

2. ADD configuration for Concurrent CHECKPOINT ENGINE

config = {
       "host_cache_size": 50,
       "parser_threads": 1,
       "pin_host_cache": True,
       "trainer": trainer,
}

3. Initialize Concurrent CHECKPOINT ENGINE

ckpt_engine = DecoratedCheckpointing(runtime_config=config, rank=args.local_rank)

4. Set Concurrent CHECKPOINT ENGINE parameter

trainer.train(resume_from_checkpoint=True)

Run Job

Execute by adding functional environment variables together with the command you want to use in the Command field of the Training Job creation screen.

| language = go

PYTHONPATH=$CHECKPOINT_VENDOR HF_DATASETS_OFFLINE="1" ${USER_SCRIPT}

• PYTHONPATH=$CHECKPOINT_VENDOR: Sets the library path to be loaded for enabling the feature.
• HF_DATASETS_OFFLINE=“1”: The Samsung Cloud Platform network does not support huggingface login or model/dataset download. * Therefore, set it to prevent huggingface network calls from user scripts.

Example

Basic code

| language = actionscript

• python file

python /mnt/experiment/training/compatiblitiy-test/version_check.py

• deepspeed

deepspeed --num_gpus=2 /mnt/experiment/training/compatiblitiy-test/train_llama_8b-demo.py

• accelerate

accelerate launch --config_file /mnt/experiment/training/compatiblitiy-test/sat-test/fsdp_config.yaml --num_processes 4 /mnt/experiment/training/compatiblitiy-test/sat-test/train_llama_1b-demo.py

When using the function

| language = actionscript

• python file

PYTHONPATH=$CHECKPOINT_VENDOR python /mnt/experiment/training/compatiblitiy-test/version_check.py

• deepspeed

PYTHONPATH=$CHECKPOINT_VENDOR deepspeed --num_gpus=4 /mnt/experiment/training/compatiblitiy-test/train_llama_8b-demo.py

• accelerate

PYTHONPATH=$CHECKPOINT_VENDOR accelerate launch --config_file /mnt/experiment/training/compatiblitiy-test/sat-test/fsdp_config.yaml --num_processes 4 /mnt/experiment/training/compatiblitiy-test/sat-test/train_llama_1b-demo.py

transformers==5.10.2
numpy==2.4.6
pybind11==3.0.4
safetensors==0.8.0
torch==2.12.1
torchvision==0.27.1
datasets==4.8.4
pytest
cuda-bindings~=13.2.0
packaging<=26.0

#--- test deepspeed version library
deepspeed==0.18.9
accelerate==1.13.0

Check progress

The progress can be viewed on the Training Job Details page’s Log tab. To check the progress, follow the steps below.

1. All Services > AI/ML > Simple AI Training Click the menu. 1. Go to the Service Home page of Simple AI Training.
2. On the Service Home page, click the Training Job menu. 2. Training Job List Navigate to the page.
3. Training Job List page, click the resource to view detailed information. 3. Go to the Training Job Details page.
4. After clicking the Log tab, check the logs. 4. You can view the logs while the environment is being prepared.

| language = actionscript | title = 

[INFO] Concurrent Checkpoint library is installed
waiting for validator through /channel/stage.socket...
validator is running....
sidecar container is running and ready for training process to run!

| language = actionscript | title = 

[2026-07-10 01:50:09,939] [ERROR] [decorator.py:442:get_last_checkpoint_preprocess] [Concurrent Checkpoint] No Checkpoint found with step: -1
ERROR:datastates.llm.decorator:[Concurrent Checkpoint] No Checkpoint found with step: -1