Concurrent Checkpointing

Concurrent Checkpointing은 기존 방식과 달리 Forward/Backward 연산 중에도 체크포인트를 비동기적으로 저장하는 기능입니다.
따라서 이 기능을 사용하여 체크포인트 저장으로 인한 오버헤드를 줄여 전체 학습 시간을 효율적으로 단축하고, 예측 불가능한 학습 중단 발생 시 유실될 수 있는 학습 진행 상태를 자동으로 저장할 수 있습니다.

Concurrent Checkpointing 개요

Simple AI Training에서 제공하는 Training Job은 On-Demand Training 타입과 Spot Training 타입이 있습니다.
Concurrent Checkpointing은 두 가지 타입에서 모두 사용할 수 있으나 타입별로 기능의 활용 범위가 다릅니다. 타입별 활용 범위는 다음과 같습니다.

Concurrent Checkpointing 사용하기

사전 준비하기: 스크립트 작성

사용자는 Trainer, Concurrent Checkpoint의 save method를 동시에 사용할 수 있습니다.

Spot Training 사용법

Spot Training 사용 시 스크립트 예시는 다음과 같습니다.

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

// 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)

// 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 저장 경로

저장 경로는 기본적으로 TrainingArgument의 output_dir을 기반으로 합니다.
해당 output_dir 하위 경로에 concurrent_checkpoint라는 디렉토리 하위에 저장됩니다.

On-Demand Training 사용하기

On-Demand Training 사용 방법 Spot Training 사용 방법과 유사합니다.

최초 학습

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

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)

수동 활성화 시

최초 학습 시 지정한 output_dir 경로에 유효한 체크포인트가 확인될 경우, 학습 재실행 시 해당 체크포인트 경로를 직접 model initialize 시 지정하여 로드합니다. 또는 학습을 재실행할 때 기존 output_dir과 동일하게 지정하고 다음과 같이 설정하여 최신 체크포인트를 자동으로 로드할 수 있습니다.

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

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)

Job 실행하기

Training Job 생성 화면의 Command 필드에 사용하고자 하는 명령어와 함께 기능 관련 환경 변수를 추가하여 실행합니다.

| language = go

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

• PYTHONPATH=$CHECKPOINT_VENDOR: 기능 활성화를 위한 라이브러리 경로를 로드하도록 설정합니다.
• HF_DATASETS_OFFLINE=“1”: Samsung Cloud Platform 망은 huggingface 로그인 및 모델/데이터셋 다운로드를 지원하지 않습니다. 따라서 사용자 스크립트에서 huggingface 네트워크 호출을 방지하는 용도로 설정합니다.

예시

기본 코드

| language = actionscript

• python 파일

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

기능 사용 시

| language = actionscript

• python 파일

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

진행 상황 확인

진행 상황은 Training Job 상세 페이지의 로그 탭에서 확인할 수 있습니다.
진행 상황을 확인하려면 다음 절차를 따르세요.

1. 모든 서비스 > AI/ML > Simple AI Training 메뉴를 클릭하세요. Simple AI Training의 Service Home 페이지로 이동합니다.
2. Service Home 페이지에서 Training Job 메뉴를 클릭하세요. Training Job 목록 페이지로 이동합니다.
3. Training Job 목록 페이지에서 상세 정보를 확인할 자원을 클릭하세요. Training Job 상세 페이지로 이동합니다.
4. 로그 탭을 클릭한 후, 로그를 확인하세요. 환경 준비 상태로 로그를 확인할 수 있습니다.

| 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