You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
To simplify the process of fine-tuning models with the LAB
method, this library provides a simple training interface.
Installing the library
To get started with the library, you must clone this repository and install it via pip.
Install the library:
pip install instructlab-training
You can then install the library for development:
pip install -e ./training
Additional NVIDIA packages
This library uses the flash-attn package as well as other packages, which rely on NVIDIA-specific CUDA tooling to be installed.
If you are using NVIDIA hardware with CUDA, you need to install the following additional dependencies.
Basic install
pip install .[cuda]
Editable install (development)
pip install -e .[cuda]
<
8000
div class="markdown-heading" dir="auto">
Using the library
You can utilize this training library by importing the necessary items.
The TrainingArgs class provides most of the customization options
for training jobs. There are a number of options you can specify, such as setting
DeepSpeed config values or running a LoRA training job instead of a full fine-tune.
TrainingArgs
Field
Description
model_path
Either a reference to a HuggingFace repo or a path to a model saved in the HuggingFace format.
data_path
A path to the .jsonl training dataset. This is expected to be in the messages format.
ckpt_output_dir
Directory where trained model checkpoints will be saved.
data_output_dir
Directory where the processed training data is stored (post filtering/tokenization/masking)
max_seq_len
The maximum sequence length to be included in the training set. Samples exceeding this length will be dropped.
max_batch_len
Maximum tokens per gpu for each batch that will be handled in a single step. Used as part of the multipack calculation. If running into out-of-memory errors, try to lower this value, but not below the max_seq_len.
num_epochs
Number of epochs to run through before stopping.
effective_batch_size
The amount of samples in a batch to see before we update the model parameters.
save_samples
Number of samples the model should see before saving a checkpoint. Consider this to be the checkpoint save frequency.
learning_rate
How fast we optimize the weights during gradient descent. Higher values may lead to unstable learning performance. It's generally recommended to have a low learning rate with a high effective batch size.
warmup_steps
The number of steps a model should go through before reaching the full learning rate. We start at 0 and linearly climb up to learning_rate.
is_padding_free
Boolean value to indicate whether or not we're training a padding-free transformer model such as Granite.
random_seed
The random seed PyTorch will use.
mock_data
Whether or not to use mock, randomly generated, data during training. For debug purposes
mock_data_len
Max length of a single mock data sample. Equivalent to max_seq_len but for mock data.
deepspeed_options
Config options to specify for the DeepSpeed optimizer.
lora
Options to specify if you intend to perform a LoRA train instead of a full fine-tune.
chat_tmpl_path
Specifies the chat template / special tokens for training.
checkpoint_at_epoch
Whether or not we should save a checkpoint at the end of each epoch.
fsdp_options
The settings for controlling FSDP when it's selected as the distributed backend.
distributed_backend
Specifies which distributed training backend to use. Supported options are "fsdp" and "deepspeed".
disable_flash_attn
Disables flash attention when set to true. This allows for training on older devices.
keep_last_checkpoint_only
Determines whether we should only keep the last checkpoint directory - the previous checkpoint directory is always overwritten. The checkpoint directory is called last_epoch.
DeepSpeedOptions
This library only currently support a few options in DeepSpeedOptions:
The default is to run with DeepSpeed, so these options only currently
allow you to customize aspects of the ZeRO stage 2 optimizer.
Field
Description
cpu_offload_optimizer
Whether or not to do CPU offloading in DeepSpeed stage 2.
cpu_offload_optimizer_ratio
Floating point between 0 & 1. Specifies the ratio of parameters updating (i.e. optimizer step) on CPU side.
cpu_offload_optimizer_pin_memory
If true, offload to page-locked CPU memory. This could boost throughput at the cost of extra memory overhead.
save_samples
The number of samples to see before saving a DeepSpeed checkpoint.
For more information about DeepSpeed, see deepspeed.ai
DeepSpeed with CPU Offloading
To use DeepSpeed with CPU offloading, you'll usually encounter an issue indicating that the optimizer needed to use the Adam optimizer on CPU doesn't exist. To resolve this, please follow the following steps:
Rebuild DeepSpeed with CPUAdam:
You'll need to rebuild DeepSpeed in order for the optimizer to be present:
# uninstall deepspeed & reinstall with the flags for installing CPUAdam
pip uninstall deepspeed
DS_BUILD_CPU_ADAM=1 DS_BUILD_UTILS=1 pip install deepspeed --no-deps
Ensure -lcurand is linked correctly:
A problem that we commonly encounter is that the -lcurand linker will not be present when
DeepSpeed recompiles. To resolve this, you will need to find the location of the libcurand.so file in your machine:
find / -name 'libcurand*.so*'2>/dev/null
If libcurand.so is not present in the /usr/lib64 directory, you'll need to add a symlink. Ex:
Like DeepSpeed, we only expose a number of parameters for you to modify with FSDP.
They are listed below:
Field
Description
cpu_offload_params
When set to true, offload parameters from the accelerator onto the CPU. This is an all-or-nothing option.
sharding_strategy
Specifies the model sharding strategy that FSDP should use. Valid options are: FULL_SHARD (ZeRO-3), HYBRID_SHARD (ZeRO-3*), SHARD_GRAD_OP (ZeRO-2), and NO_SHARD.
Note
For sharding_strategy - Only SHARD_GRAD_OP has been extensively tested and is actively supported by this library.
loraOptions
LoRA options currently supported:
Field
Description
rank
The rank parameter for LoRA training.
alpha
The alpha parameter for LoRA training.
dropout
The dropout rate for LoRA training.
target_modules
The list of target modules for LoRA training.
quantize_data_type
The data type for quantization in LoRA training. Valid options are None and "nf4"
Example run with LoRa options
If you'd like to do a LoRA train, you can specify a LoRA
option to TrainingArgs via the LoraOptions object.
When running the training script, we always invoke torchrun.
If you are running a single-GPU system or something that doesn't
otherwise require distributed training configuration, you can create a default object:
However, if you want to specify a more complex configuration,
the library currently supports all the options that torchrun accepts
today.
Note
For more information about the torchrun arguments, please consult the torchrun documentation.
Example training run with TorchrunArgs arguments
For example, in a 8-GPU, 2-machine system, we would
specify the following torchrun config:
MASTER_ADDR=os.getenv('MASTER_ADDR')
MASTER_PORT=os.getnev('MASTER_PORT')
RDZV_ENDPOINT=f'{MASTER_ADDR}:{MASTER_PORT}'# on machine 1torchrun_args=TorchrunArgs(
nnodes=2, # number of machines nproc_per_node=4, # num GPUs per machinenode_rank=0, # node rank for this machinerdzv_id=123,
rdzv_endpoint=RDZV_ENDPOINT
)
run_training(
torchrun_args=torchrun_args,
training_args=training_args
)
MASTER_ADDR=os.getenv('MASTER_ADDR')
MASTER_PORT=os.getnev('MASTER_PORT')
RDZV_ENDPOINT=f'{MASTER_ADDR}:{MASTER_PORT}'# on machine 2torchrun_args=TorchrunArgs(
nnodes=2, # number of machines nproc_per_node=4, # num GPUs per machinenode_rank=1, # node rank for this machinerdzv_id=123,
rdzv_endpoint=f'{MASTER_ADDR}:{MASTER_PORT}'
)
run_training(
torch_args=torchrun_args,
train_args=training_args
)
Example training run with arguments
Define the training arguments which will serve as the
parameters for our training run:
# define training-specific argumentstraining_args=TrainingArgs(
# define data-specific argumentsmodel_path="ibm-granite/granite-7b-base",
data_path="path/to/dataset.jsonl",
ckpt_output_dir="data/saved_checkpoints",
data_output_dir="data/outputs",
# define model-trianing parametersmax_seq_len=4096,
max_batch_len=60000,
num_epochs=10,
effective_batch_size=3840,
save_samples=250000,
learning_rate=2e-6,
warmup_steps=800,
is_padding_free=True, # set this to true when using Granite-based modelsrandom_seed=42,
)
We'll also need to define the settings for running a multi-process job
via torchrun. To do this, create a TorchrunArgs object.
Note, for single-GPU jobs, you can simply set nnodes = 1 and nproc_per_node=1.
torchrun_args=TorchrunArgs(
nnodes=1, # number of machines nproc_per_node=8, # num GPUs per machinenode_rank=0, # node rank for this machinerdzv_id=123,
rdzv_endpoint='127.0.0.1:12345'
)
Finally, you can just call run_training and this library will handle
the rest 🙂.
Example training with separate data pre-processing
If the machines in the example above have shared storage, users can pre-process the training dataset a single time so that it can then be distributed to each machine by making the following updates.
