hanzo-ml/trainer
1
0
Fork 0
mirror of https://github.com/hanzo-ml/trainer.git synced 2026-09-24 00:31:34 +00:00
Code Issues Packages Projects Releases Wiki Activity
Files
release-2.2
trainer/examples/deepspeed/text-summarization/T5-Fine-Tuning.ipynb
T
Andrey VelichkevichandGitHub fbacd9f577 fix(examples): Verify TrainJob Completion (#3331) 
* fix(examples): Verify TrainJob Completion

Signed-off-by: Andrey Velichkevich <[email protected]>

* Rename client to TrainerClient()

Signed-off-by: Andrey Velichkevich <[email protected]>

---------

Signed-off-by: Andrey Velichkevich <[email protected]>
2026-03-16 14:40:31 +00:00

64 KiB
Raw Permalink Blame History

T5 LLM Fine-Tuning with DeepSpeed and Kubeflow Trainer

This Notebook will fine-tune Text-to-Text Transfer Transformer (T5) with CNN Dailymail dataset for text summarization using Kubeflow TrainJob and DeepSpeed.

Pretrained T5 model: https://huggingface.co/google-t5/t5-base

CNN Dailymail dataset: https://huggingface.co/datasets/abisee/cnn_dailymail

This Notebook will use 4 x A100 NVIDIA GPUs, to fine-tune T5 model on 2 nodes (every node has 4 GPUs).

Install the Kubeflow SDK

You need to install the Kubeflow SDK to interact with Kubeflow Trainer APIs:

In [1]:
# !pip install -U kubeflow

Create Script to Fine-Tune T5 with DeepSpeed

We need to wrap our fine-tuning script into a function to create Kubeflow TrainJob.

In [ ]:
def deepspeed_train_t5(num_samples: str, model_name: str, bucket: str):
    import os
    import time
    import boto3
    import torch
    import torch.distributed as dist
    from torch.utils.data.distributed import DistributedSampler
    from transformers import T5Tokenizer, T5ForConditionalGeneration
    from datasets import load_dataset
    import deepspeed
    import numpy as np

    # Initialize distributed environment.
    deepspeed.init_distributed(dist_backend="nccl")
    local_rank = int(os.environ["LOCAL_RANK"])

    # Define the summarization dataset class
    class SummarizationDataset(torch.utils.data.Dataset):
        def __init__(self, tokenizer, num_samples):
            # Load publicly available CNN/DailyMail dataset
            self.dataset = load_dataset("cnn_dailymail", "3.0.0", split=f"train[:{num_samples}]")
            self.tokenizer = tokenizer

        def __len__(self):
            return len(self.dataset)

        def clean_text(self, text):
            if text is None:
                return ""
            return text.replace("\n", " ").replace("``", "").replace('"', "").strip()

        def convert_to_features(self, example_batch):
            # Map CNN/DailyMail fields to 'text' and 'summary'
            input_ = self.clean_text(example_batch["article"])
            target_ = self.clean_text(example_batch["highlights"])

            source = self.tokenizer(
                input_,
                max_length=512,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )
            targets = self.tokenizer(
                target_,
                max_length=150,
                padding="max_length",
                truncation=True,
                return_tensors="pt",
            )

            return source, targets

        def __getitem__(self, index):
            source, targets = self.convert_to_features(self.dataset[index])
            return {
                "source_ids": source["input_ids"].squeeze(),
                "source_mask": source["attention_mask"].squeeze(),
                "target_ids": targets["input_ids"].squeeze(),
                "target_mask": targets["attention_mask"].squeeze(),
            }

    # Download model and tokenizer.
    if dist.get_rank() == 0:
        print("-" * 100)
        print("Downloading T5 Model")
        print("-" * 100)

    model = T5ForConditionalGeneration.from_pretrained(model_name)
    tokenizer = T5Tokenizer.from_pretrained(model_name)

    # Download dataset.
    dataset = SummarizationDataset(tokenizer, num_samples=int(num_samples))
    train_loader = torch.utils.data.DataLoader(
        dataset, batch_size=4, sampler=DistributedSampler(dataset)
    )

    # Define DeepSpeed configuration.
    # Train batch size = micro batch size * gradient steps * GPUs (e.g. 2 x 1 x 8 = 16).
    ds_config = {
        "train_micro_batch_size_per_gpu": 2,
        "gradient_accumulation_steps": 1,
        # "fp16": {"enabled": True}, # If your GPU (e.g. V100) doesn't support bf16, use fp16.
        "bf16": {"enabled": True}, # Enable mixed precision.
        "optimizer": {
            "type": "AdamW",
            "params": {"lr": 0.002},
        },
        "scheduler": {
            "type": "WarmupLR",
            "params": {
                "warmup_min_lr": 0,
                "warmup_max_lr": 0.001,
                "warmup_num_steps": 1000,
            },
        },
    }

    # Initialize model with DeepSpeed.
    model, _, _, _ = deepspeed.initialize(
        config=ds_config,
        model=model,
        model_parameters=model.parameters(),
    )

    # Start training process.
    if dist.get_rank() == 0:
        print("-" * 100)
        print("Starting DeepSpeed distributed training...")
        print("-" * 100)

    t0 = time.time()
    for epoch in range(1, 3):
        losses = []
        for batch_idx, batch in enumerate(train_loader):
            for key in batch.keys():
                batch[key] = batch[key].to(local_rank)
            # Forward pass.
            output = model(
                input_ids=batch["source_ids"],
                attention_mask=batch["source_mask"],
                labels=batch["target_ids"],
            )
            loss = output.loss

            # Run backpropagation.
            model.backward(loss)
            # Weight updates.
            model.step()
            losses.append(loss.item())

            if batch_idx % 10 == 0 and dist.get_rank() == 0:
                print(
                    f"Train Epoch: {epoch} [{batch_idx * len(batch)}/{len(train_loader.dataset)} "
                    f"({100.0 * batch_idx / len(train_loader):.0f}%)]\tLoss: {loss.item():.6f}"
                )

        if dist.get_rank() == 0:
            print("-" * 100)
            print(f"Average Train Loss: {np.mean(losses):.4f}")
            print("-" * 100)

    if dist.get_rank() == 0:
        print("-" * 100)
        print(f"DeepSpeed training time: {int(time.time() - t0)} seconds")
        print("-" * 100)

        print("Exporting HuggingFace model to S3")
        MODEL_PATH = os.path.join("/home/mpiuser", model_name)
        model.module.save_pretrained(MODEL_PATH)
        tokenizer.save_pretrained(MODEL_PATH)

        s3_bucket = boto3.resource("s3").Bucket(bucket)
        for file in os.listdir(MODEL_PATH):
            print(f"Uploading file {os.path.join(MODEL_PATH, file)}")
            s3_bucket.upload_file(
                os.path.join(MODEL_PATH, file), os.path.join(model_name, file)
            )

List Available Kubeflow Trainer Runtimes

Get available Kubeflow Trainer Runtimes with the list_runtimes() API.

You can inspect Runtime details, including the name, framework, and available default devices on the single node.

  • Runtimes with CustomTrainer: You must write the training script within the function.

  • Runtimes with BuiltinTrainer: You can configure settings (e.g., LoRA Config) for LLM fine-tuning Job.

In [3]:
from kubeflow.trainer import TrainerClient, CustomTrainer

for r in TrainerClient().list_runtimes():
    if r.name == "deepspeed-distributed":
        print(f"Name: {r.name}, Framework: {r.trainer.framework}, Trainer Type: {r.trainer.trainer_type.value}\n")
        print(f"Runtime devices: {r.trainer.device} x {r.trainer.device_count}")
        deepspeed_runtime = r
Name: deepspeed-distributed, Framework: deepspeed, Trainer Type: CustomTrainer

Runtime devices: gpu x 1

Create TrainJob for Distributed Training

Use the train() API to scale the training code across 2 Nodes and 8 GPUs.

Don't forget to update the S3 bucket name.

In [9]:
MODEL_NAME = "t5-base"
# BUCKET_NAME = "TODO: add your bucket here"
In [10]:
args = {
    "num_samples": "2000",
    "model_name": MODEL_NAME,
    "bucket": BUCKET_NAME,
}

job_id = TrainerClient().train(
    trainer=CustomTrainer(
        func=deepspeed_train_t5,
        func_args=args,
        packages_to_install=["boto3"], # Custom packages to install at runtime.
        num_nodes=2,
        resources_per_node={
            "gpu": 4
        },
    ),
    runtime=deepspeed_runtime,
)
In [11]:
# Train API generates a random TrainJob id.
job_id
Out [11]:
'l82624f2928b'

Check the TrainJob Info

Use the list_jobs() and get_job() APIs to get information about created TrainJob and its steps.

In [12]:
for job in TrainerClient().list_jobs():
    print(f"TrainJob: {job.name}, Status: {job.status}, Created at: {job.creation_timestamp}")
TrainJob: l82624f2928b, Status: Created, Created at: 2025-10-18 00:07:48+00:00
In [14]:
# We execute mpirun command on node-0, which functions as the MPI Launcher node.
for c in TrainerClient().get_job(name=job_id).steps:
    print(f"Step: {c.name}, Status: {c.status}, Devices: {c.device} x {c.device_count}\n")
Step: node-0, Status: Running, Devices: gpu x 4

Step: node-1, Status: Running, Devices: gpu x 4

Get the TrainJob Logs

Use the get_job_logs() API to retrieve the TrainJob logs.

Since we distribute the dataset accross 8 GPUs (2 nodes x 4 GPUs), each rank processes round(2000 / 8) = 250 samples.

In [15]:
for logline in TrainerClient().get_job_logs(job_id, follow=True):
    print(logline)
Warning: Permanently added '[l82624f2928b-node-0-0.l82624f2928b]:2222' (ECDSA) to the list of known hosts.
[2025-10-18 00:08:23,111] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
df: /home/mpiuser/.triton/autotune: No such file or directory
[2025-10-18 00:08:24,295] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:24,754] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:24,767] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:24,767] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:25,593] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:25,910] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:25,923] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:25,923] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:27,422] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:27,435] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:27,435] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:30,514] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:31,987] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:32,000] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:32,000] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:40,410] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
df: /home/mpiuser/.triton/autotune: No such file or directory
[2025-10-18 00:08:41,794] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:42,072] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:42,085] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:42,085] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:43,484] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:43,497] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:43,498] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:49,499] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:50,045] [INFO] [real_accelerator.py:254:get_accelerator] Setting ds_accelerator to cuda (auto detect)
[2025-10-18 00:08:50,997] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:51,010] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:51,010] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:51,680] [INFO] [logging.py:107:log_dist] [Rank -1] [TorchCheckpointEngine] Initialized with serialization = False
[2025-10-18 00:08:51,692] [INFO] [comm.py:821:init_distributed] cdb=None
[2025-10-18 00:08:51,693] [INFO] [comm.py:836:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=0, local_rank=0, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=2, local_rank=2, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=1, local_rank=1, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=3, local_rank=3, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:852:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=4, local_rank=0, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=6, local_rank=2, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=5, local_rank=1, world_size=8, master_addr=240.56.142.151, master_port=29500
[2025-10-18 00:08:51,982] [INFO] [comm.py:891:mpi_discovery] Discovered MPI settings of world_rank=7, local_rank=3, world_size=8, master_addr=240.56.142.151, master_port=29500
----------------------------------------------------------------------------------------------------
Downloading T5 Model
----------------------------------------------------------------------------------------------------
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
You are using the default legacy behaviour of the <class 'transformers.models.t5.tokenization_t5.T5Tokenizer'>. This is expected, and simply means that the `legacy` (previous) behavior will be used so nothing changes for you. If you want to use the new behaviour, set `legacy=False`. This should only be set if you understand what it means, and thoroughly read the reason why this was added as explained in https://github.com/huggingface/transformers/pull/24565
Generating train split: 100%|██████████| 287113/287113 [00:03<00:00, 77899.43 examples/s]
Generating train split: 100%|██████████| 287113/287113 [00:03<00:00, 79219.92 examples/s]s]
Generating validation split: 100%|██████████| 13368/13368 [00:00<00:00, 74240.68 examples/s]
Generating validation split: 100%|██████████| 13368/13368 [00:00<00:00, 80800.92 examples/s]
Generating test split: 100%|██████████| 11490/11490 [00:00<00:00, 81368.13 examples/s]
[2025-10-18 00:09:10,600] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,609] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed info: version=0.17.4, git-hash=unknown, git-branch=unknown
[2025-10-18 00:09:10,609] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,610] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,613] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
Generating test split: 100%|██████████| 11490/11490 [00:00<00:00, 82023.46 examples/s]
[2025-10-18 00:09:10,728] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,738] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,746] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:10,758] [INFO] [config.py:684:__init__] Config mesh_device None world_size = 8
[2025-10-18 00:09:11,696] [INFO] [engine.py:1339:_configure_distributed_model] ********** distributed groups summary **********
	 self.dp_world_size=8
	 self.mp_world_size=1
	 self.seq_dp_world_size=8
	 self.sequence_parallel_size=1
***********************************************
[2025-10-18 00:09:11,843] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Creating extension directory /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Creating extension directory /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Using /home/mpiuser/.cache/torch_extensions/py310_cu126 as PyTorch extensions root...
Detected CUDA files, patching ldflags
Emitting ninja build file /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam/build.ninja...
Detected CUDA files, patching ldflags
Emitting ninja build file /home/mpiuser/.cache/torch_extensions/py310_cu126/fused_adam/build.ninja...
/home/mpiuser/.local/lib/python3.10/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. 
If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
  warnings.warn(
/home/mpiuser/.local/lib/python3.10/site-packages/torch/utils/cpp_extension.py:2356: UserWarning: TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. 
If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'].
  warnings.warn(
Building extension module fused_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
Building extension module fused_adam...
Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)
[1/3] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output multi_tensor_adam.cuda.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80 -UC10_USE_GLOG -DBF16_AVAILABLE -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -std=c++17 -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o 
[1/3] /usr/local/cuda/bin/nvcc --generate-dependencies-with-compile --dependency-output multi_tensor_adam.cuda.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr -gencode=arch=compute_80,code=compute_80 -gencode=arch=compute_80,code=sm_80 --compiler-options '-fPIC' -O3 -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -lineinfo --use_fast_math -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_80,code=compute_80 -UC10_USE_GLOG -DBF16_AVAILABLE -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -std=c++17 -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/multi_tensor_adam.cu -o multi_tensor_adam.cuda.o 
[2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -UC10_USE_GLOG -DBF16_AVAILABLE -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o 
[2/3] c++ -MMD -MF fused_adam_frontend.o.d -DTORCH_EXTENSION_NAME=fused_adam -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE=\"_gcc\" -DPYBIND11_STDLIB=\"_libstdcpp\" -DPYBIND11_BUILD_ABI=\"_cxxabi1016\" -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/includes -I/home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include -isystem /home/mpiuser/.local/lib/python3.10/site-packages/torch/include/torch/csrc/api/include -isystem /usr/local/cuda/include -isystem /usr/include/python3.10 -D_GLIBCXX_USE_CXX11_ABI=1 -fPIC -std=c++17 -O3 -std=c++17 -g -Wno-reorder -DVERSION_GE_1_1 -DVERSION_GE_1_3 -DVERSION_GE_1_5 -UC10_USE_GLOG -DBF16_AVAILABLE -c /home/mpiuser/.local/lib/python3.10/site-packages/deepspeed/ops/csrc/adam/fused_adam_frontend.cpp -o fused_adam_frontend.o 
[3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/home/mpiuser/.local/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o fused_adam.so
Loading extension module fused_adam...
Time to load fused_adam op: 25.939489603042603 seconds
[2025-10-18 00:09:37,813] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
Loading extension module fused_adam...
Time to load fused_adam op: 26.032996654510498 seconds
Loading extension module fused_adam...
Time to load fused_adam op: 26.032512187957764 seconds
Loading extension module fused_adam...
Time to load fused_adam op: 26.037198543548584 seconds
[2025-10-18 00:09:37,883] [INFO] [logging.py:107:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer
[2025-10-18 00:09:37,883] [INFO] [logging.py:107:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2025-10-18 00:09:37,894] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Basic Optimizer = FusedAdam
[2025-10-18 00:09:37,894] [INFO] [logging.py:107:log_dist] [Rank 0] Creating BF16 optimizer
[3/3] c++ fused_adam_frontend.o multi_tensor_adam.cuda.o -shared -L/home/mpiuser/.local/lib/python3.10/site-packages/torch/lib -lc10 -lc10_cuda -ltorch_cpu -ltorch_cuda -ltorch -ltorch_python -L/usr/local/cuda/lib64 -lcudart -o fused_adam.so
Loading extension module fused_adam...
Time to load fused_adam op: 26.056629180908203 seconds
[2025-10-18 00:09:37,908] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:37,915] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:37,954] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
Loading extension module fused_adam...
Loading extension module fused_adam...
Time to load fused_adam op: 26.13313388824463 seconds
Time to load fused_adam op: 26.13256287574768 seconds
Loading extension module fused_adam...
Time to load fused_adam op: 26.137150526046753 seconds
[2025-10-18 00:09:38,010] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:38,011] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:38,018] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:38,120] [INFO] [utils.py:781:see_memory_usage] begin bf16_optimizer
[2025-10-18 00:09:38,121] [INFO] [utils.py:782:see_memory_usage] MA 0.43 GB         Max_MA 0.43 GB         CA 0.44 GB         Max_CA 0 GB 
[2025-10-18 00:09:38,121] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 23.24 GB, percent = 2.1%
[2025-10-18 00:09:38,314] [INFO] [utils.py:781:see_memory_usage] before initializing group 0
[2025-10-18 00:09:38,315] [INFO] [utils.py:782:see_memory_usage] MA 0.43 GB         Max_MA 0.43 GB         CA 0.44 GB         Max_CA 0 GB 
[2025-10-18 00:09:38,315] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 24.34 GB, percent = 2.2%
[2025-10-18 00:09:38,523] [INFO] [utils.py:781:see_memory_usage] after initializing group 0
[2025-10-18 00:09:38,524] [INFO] [utils.py:782:see_memory_usage] MA 1.35 GB         Max_MA 1.35 GB         CA 1.85 GB         Max_CA 2 GB 
[2025-10-18 00:09:38,524] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 25.03 GB, percent = 2.2%
[2025-10-18 00:09:38,704] [INFO] [utils.py:781:see_memory_usage] end bf16_ optimizer
[2025-10-18 00:09:38,705] [INFO] [utils.py:782:see_memory_usage] MA 1.35 GB         Max_MA 1.35 GB         CA 1.85 GB         Max_CA 2 GB 
[2025-10-18 00:09:38,705] [INFO] [utils.py:789:see_memory_usage] CPU Virtual Memory:  used = 25.14 GB, percent = 2.2%
[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed Final Optimizer = BF16_Optimizer
[2025-10-18 00:09:38,705] [WARNING] [lr_schedules.py:686:get_lr] Attempting to get learning rate from scheduler before it has started
[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = WarmupLR
[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupLR object at 0x7f91ac245fc0>
[2025-10-18 00:09:38,705] [INFO] [logging.py:107:log_dist] [Rank 0] step=0, skipped=0, lr=[0], mom=[(0.9, 0.999)]
[2025-10-18 00:09:38,706] [INFO] [logging.py:107:log_dist] [Rank 0] [TorchCheckpointEngine] Initialized with serialization = True
[2025-10-18 00:09:38,706] [INFO] [config.py:954:print] DeepSpeedEngine configuration:
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   activation_checkpointing_config  {
    "partition_activations": false, 
    "contiguous_memory_optimization": false, 
    "cpu_checkpointing": false, 
    "number_checkpoints": null, 
    "synchronize_checkpoint_boundary": false, 
    "profile": false
}
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'intra_op_parallelism': 1, 'single_submit': False, 'overlap_events': True, 'use_gds': False}
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   amp_enabled .................. False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   amp_params ................... False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   autotuning_config ............ {
    "enabled": false, 
    "start_step": null, 
    "end_step": null, 
    "metric_path": null, 
    "arg_mappings": null, 
    "metric": "throughput", 
    "model_info": null, 
    "results_dir": "autotuning_results", 
    "exps_dir": "autotuning_exps", 
    "overwrite": true, 
    "fast": true, 
    "start_profile_step": 3, 
    "end_profile_step": 5, 
    "tuner_type": "gridsearch", 
    "tuner_early_stopping": 5, 
    "tuner_num_trials": 50, 
    "model_info_path": null, 
    "mp_size": 1, 
    "max_train_batch_size": null, 
    "min_train_batch_size": 1, 
    "max_train_micro_batch_size_per_gpu": 1.024000e+03, 
    "min_train_micro_batch_size_per_gpu": 1, 
    "num_tuning_micro_batch_sizes": 3
}
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   bfloat16_config .............. enabled=True immediate_grad_update=False check_grad_overflow=False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_config ............ {'tag_validation': 'WARN', 'checkpoint_serialization': True, 'writer': None}
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_parallel_write_pipeline  False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_tag_validation_enabled  True
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   checkpoint_tag_validation_fail  False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7f91ac289630>
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   communication_data_type ...... None
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   compile_config ............... deepcompile=False free_activation=False offload_activation=False offload_opt_states=False double_buffer=True symmetric_memory=False debug_log=False offload_parameters=False sync_before_reduce=False sync_after_reduce=False sync_before_allgather=False sync_after_allgather=False keep_int_input_tensors=True keep_all_input_tensors=False
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}
[2025-10-18 00:09:38,707] [INFO] [config.py:958:print]   curriculum_enabled_legacy .... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   curriculum_params_legacy ..... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'pin_memory': False, 'curriculum_learning': {'enabled': False}, 'dynamic_batching': {'enabled': False, 'lr_scaling_method': 'linear', 'min_batch_size': 1, 'max_batch_size': None, 'sequence_picking_order': 'dataloader', 'verbose': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   data_efficiency_enabled ...... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   dataloader_drop_last ......... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   disable_allgather ............ False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   dump_state ................... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_enabled ........... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_gas_boundary_resolution  1
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_layer_name ........ bert.encoder.layer
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_layer_num ......... 0
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_max_iter .......... 100
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_stability ......... 1e-06
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_tol ............... 0.01
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   eigenvalue_verbose ........... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   elasticity_enabled ........... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   float16_config ............... enabled=False auto_cast=False loss_scale=0.0 initial_scale_power=16 loss_scale_window=1000 hysteresis=2 consecutive_hysteresis=False min_loss_scale=1 fp16_master_weights_and_grads=False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   flops_profiler_config ........ {
    "enabled": false, 
    "recompute_fwd_factor": 0.0, 
    "profile_step": 1, 
    "module_depth": -1, 
    "top_modules": 1, 
    "detailed": true, 
    "output_file": null
}
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   global_rank .................. 0
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   grad_accum_dtype ............. None
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_accumulation_steps .. 1
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_clipping ............ 0.0
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   gradient_predivide_factor .... 1.0
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   graph_harvesting ............. False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   load_universal_checkpoint .... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   memory_breakdown ............. False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   mics_hierarchial_params_gather  False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   mics_shard_size .............. -1
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') comet=CometConfig(enabled=False, samples_log_interval=100, project=None, workspace=None, api_key=None, experiment_name=None, experiment_key=None, online=None, mode=None) wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName')
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   nebula_config ................ {
    "enabled": false, 
    "persistent_storage_path": null, 
    "persistent_time_interval": 100, 
    "num_of_version_in_retention": 2, 
    "enable_nebula_load": true, 
    "load_path": null
}
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_legacy_fusion ...... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_name ............... adamw
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   optimizer_params ............. {'lr': 0.002}
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0, 'pipe_partitioned': True, 'grad_partitioned': True}
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pld_enabled .................. False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   pld_params ................... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   prescale_gradients ........... False
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   scheduler_name ............... WarmupLR
[2025-10-18 00:09:38,708] [INFO] [config.py:958:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 0.001, 'warmup_num_steps': 1000}
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   seq_parallel_communication_data_type  torch.float32
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   sparse_attention ............. None
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   sparse_gradients_enabled ..... False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   steps_per_print .............. None
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   tensor_parallel_config ....... dtype=torch.float16 autotp_size=0 tp_overlap_comm=False tensor_parallel=TPConfig(tp_size=1, tp_grain_size=1, mpu=None, tp_group=None) injection_policy_tuple=None keep_module_on_host=False replace_with_kernel_inject=False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   timers_config ................ enabled=True synchronized=True
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_dtype ......... None
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_enabled ....... False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   torch_autocast_lower_precision_safe_modules  None
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   train_batch_size ............. 16
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   train_micro_batch_size_per_gpu  2
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   use_data_before_expert_parallel_  False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   use_node_local_storage ....... False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   wall_clock_breakdown ......... False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   weight_quantization_config ... None
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   world_size ................... 8
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_allow_untested_optimizer  False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_config .................. stage=0 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=500000000 use_multi_rank_bucket_allreduce=True allgather_partitions=True allgather_bucket_size=500000000 overlap_comm=False load_from_fp32_weights=True elastic_checkpoint=False offload_param=None offload_optimizer=None sub_group_size=1000000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=50000000 param_persistence_threshold=100000 model_persistence_threshold=9223372036854775807 max_live_parameters=1000000000 max_reuse_distance=1000000000 gather_16bit_weights_on_model_save=False module_granularity_threshold=0 use_all_reduce_for_fetch_params=False stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False zeropp_loco_param=None mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True log_trace_cache_warnings=False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_enabled ................. False
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_force_ds_cpu_optimizer .. True
[2025-10-18 00:09:38,709] [INFO] [config.py:958:print]   zero_optimization_stage ...... 0
[2025-10-18 00:09:38,709] [INFO] [config.py:944:print_user_config]   json = {
    "train_micro_batch_size_per_gpu": 2, 
    "gradient_accumulation_steps": 1, 
    "bf16": {
        "enabled": true
    }, 
    "optimizer": {
        "type": "AdamW", 
        "params": {
            "lr": 0.002
        }
    }, 
    "scheduler": {
        "type": "WarmupLR", 
        "params": {
            "warmup_min_lr": 0, 
            "warmup_max_lr": 0.001, 
            "warmup_num_steps": 1000
        }
    }
}
----------------------------------------------------------------------------------------------------
Starting DeepSpeed distributed training...
----------------------------------------------------------------------------------------------------
Train Epoch: 1 [0/2000 (0%)]	Loss: 10.250000
Train Epoch: 1 [40/2000 (16%)]	Loss: 0.878906
Train Epoch: 1 [80/2000 (32%)]	Loss: 0.960938
Train Epoch: 1 [120/2000 (48%)]	Loss: 0.957031
Train Epoch: 1 [160/2000 (63%)]	Loss: 0.707031
Train Epoch: 1 [200/2000 (79%)]	Loss: 0.621094
Train Epoch: 1 [240/2000 (95%)]	Loss: 0.761719
----------------------------------------------------------------------------------------------------
Average Train Loss: 1.5145
----------------------------------------------------------------------------------------------------
Train Epoch: 2 [0/2000 (0%)]	Loss: 0.625000
Train Epoch: 2 [40/2000 (16%)]	Loss: 0.558594
Train Epoch: 2 [80/2000 (32%)]	Loss: 0.617188
Train Epoch: 2 [120/2000 (48%)]	Loss: 0.667969
Train Epoch: 2 [160/2000 (63%)]	Loss: 0.494141
Train Epoch: 2 [200/2000 (79%)]	Loss: 0.365234
Train Epoch: 2 [240/2000 (95%)]	Loss: 0.523438
----------------------------------------------------------------------------------------------------
Average Train Loss: 0.5215
----------------------------------------------------------------------------------------------------
----------------------------------------------------------------------------------------------------
DeepSpeed training time: 59 seconds
----------------------------------------------------------------------------------------------------
Exporting HuggingFace model to S3
Uploading file /home/mpiuser/t5-base/config.json
Uploading file /home/mpiuser/t5-base/generation_config.json
Uploading file /home/mpiuser/t5-base/model.safetensors
[rank4]:[W1018 00:10:39.692111139 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
Uploading file /home/mpiuser/t5-base/tokenizer_config.json
Uploading file /home/mpiuser/t5-base/special_tokens_map.json
Uploading file /home/mpiuser/t5-base/added_tokens.json
Uploading file /home/mpiuser/t5-base/spiece.model
[rank0]:[W1018 00:10:44.046274646 ProcessGroupNCCL.cpp:1479] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())

Verify TrainJob Completion

Check the TrainJob status to ensure it completed successfully.

In [ ]:
TrainerClient().wait_for_job_status(name=job_id, timeout=20)

Download the Trained Model

Finally, download fine-tuned model from S3 for evaluations.

In [20]:
LOCAL_DIR = "./t5-base"
In [21]:
import boto3
import os

os.makedirs(LOCAL_DIR, exist_ok=True)
s3 = boto3.client("s3")
for obj in s3.list_objects_v2(Bucket=BUCKET_NAME, Prefix="t5-base")["Contents"]:
    file = obj["Key"]

    print(f"Downloading file: {file}")
    s3.download_file(BUCKET_NAME, file, os.path.join(LOCAL_DIR, os.path.basename(file)))
Downloading file: t5-base/added_tokens.json
Downloading file: t5-base/config.json
Downloading file: t5-base/generation_config.json
Downloading file: t5-base/model.safetensors
Downloading file: t5-base/special_tokens_map.json
Downloading file: t5-base/spiece.model
Downloading file: t5-base/tokenizer_config.json

Evaluate Fine-Tuned T5 Model

After model is downloaded, you can load it into the HuggingFace pipeline.

The T5 model performs well for NLP tasks such as summarization, translation, and text classification.

In the example below, we'll demonstrate how to use a fine-tuned version of the T5 model to summarize documentation related to the Kubeflow Trainer project.

In [24]:
import torch
from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, pipeline

# Load the fine-tuned T5 model.
model = AutoModelForSeq2SeqLM.from_pretrained(LOCAL_DIR)
tokenizer = AutoTokenizer.from_pretrained(LOCAL_DIR)

summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, framework="pt")

text = """
summarize: In Kubeflow Trainer you can integrate other ML libraries such as HuggingFace,
DeepSpeed, or Megatron-LM with Kubeflow Trainer to orchestrate their ML training on Kubernetes.
Kubeflow Trainer allows you to effortlessly develop your LLMs with the Kubeflow Python SDK
and build Kubernetes-native Training Runtimes with Kubernetes Custom Resources APIs.
Kubeflow Trainer is a Kubernetes-native project designed for large language models (LLMs)
fine-tuning and enabling scalable, distributed training of machine learning (ML)
models across various frameworks, including PyTorch, JAX, TensorFlow, and XGBoost.
"""

summarizer(text, min_length=5, max_length=100)
Out [24]:
Device set to use mps:0
Both `max_new_tokens` (=256) and `max_length`(=100) seem to have been set. `max_new_tokens` will take precedence. Please refer to the documentation for more information. (https://huggingface.co/docs/transformers/main/en/main_classes/text_generation)

[{'summary_text': 'Kubeflow Trainer allows you to fine-tune and enabling scalable, distributed training of machine learning (ML) models . Kubernetes-native project designed for large language models (LLMs).'}]
In [ ]:
Reference in New Issue View Git Blame Copy Permalink