kaira.training.TrainingArguments

Inheritance diagram of TrainingArguments — Inheritance diagram for TrainingArguments

class kaira.training.TrainingArguments(snr_min: float = 0.0, snr_max: float = 20.0, noise_variance_min: float = 0.1, noise_variance_max: float = 2.0, channel_uses: int | None = None, code_length: int | None = None, info_length: int | None = None, channel_type: str = 'awgn', push_to_hub: bool = False, hub_model_id: str | None = None, hub_token: str | None = None, hub_private: bool = False, hub_strategy: str = 'end', output_dir: str = './results', num_train_epochs: float = 10.0, per_device_train_batch_size: int = 32, per_device_eval_batch_size: int = 32, learning_rate: float = 0.0001, warmup_steps: int = 1000, logging_steps: int = 100, eval_steps: int = 500, save_steps: int = 1000, eval_strategy: str = 'steps', logging_strategy: str = 'steps', save_strategy: str = 'steps', **kwargs)[source]

Bases: TrainingArguments

Training arguments that support Hydra configuration management.

This class extends transformers.TrainingArguments to provide seamless integration with Hydra configuration management while maintaining full compatibility with Hugging Face ecosystem. It supports:

Direct instantiation from Hydra DictConfig via from_hydra_config
Communication-specific parameters
Automatic parameter filtering and validation

Examples

>>> # From Hydra config
>>> hydra_config = OmegaConf.create({"training": {"output_dir": "./results", "num_train_epochs": 10}})
>>> args = TrainingArguments.from_hydra_config(hydra_config)

>>> # With communication parameters
>>> args = TrainingArguments(
...     output_dir="./results",
...     snr_min=0.0,
...     snr_max=20.0,
...     channel_uses=64
... )

Methods

`__init__`	Initialize TrainingArguments.
`from_cli_args`	Create TrainingArguments from command-line arguments.
`from_hydra_config`	Create TrainingArguments from Hydra configuration.
`get_noise_variance_range`	Get noise variance range as tuple.
`get_process_log_level`	Returns the log level to be used depending on whether this process is the main process of node 0, main process of node non-0, or a non-main process.
`get_snr_range`	Get SNR range as tuple.
`get_warmup_steps`	Get number of steps used for a linear warmup.
`main_process_first`	A context manager for torch distributed environment where on needs to do something on the main process, while blocking replicas, and when it's finished releasing the replicas.
`to_dict`	Convert to dictionary representation.
`to_hydra_config`	Convert to Hydra DictConfig.
`to_json_string`	Serializes this instance to a JSON string.
`to_sanitized_dict`	Sanitized serialization to use with TensorBoard’s hparams

Attributes

`accelerator_config`
`adafactor`
`adam_beta1`
`adam_beta2`
`adam_epsilon`
`auto_find_batch_size`
`average_tokens_across_devices`
`batch_eval_metrics`
`bf16`
`bf16_full_eval`
`data_seed`
`dataloader_drop_last`
`dataloader_num_workers`
`dataloader_persistent_workers`
`dataloader_pin_memory`
`dataloader_prefetch_factor`
`ddp_backend`
`ddp_broadcast_buffers`
`ddp_bucket_cap_mb`
`ddp_find_unused_parameters`
`ddp_timeout`
`ddp_timeout_delta`	The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable.
`debug`
`deepspeed`
`default_optim`
`device`	The device used by this process.
`disable_tqdm`
`do_eval`
`do_predict`
`do_train`
`eval_accumulation_steps`
`eval_batch_size`	The actual batch size for evaluation (may differ from per_gpu_eval_batch_size in distributed training).
`eval_delay`
`eval_do_concat_batches`
`eval_on_start`
`eval_steps`
`eval_strategy`
`eval_use_gather_object`
`fp16`
`fp16_backend`
`fp16_full_eval`
`fp16_opt_level`
`framework`
`fsdp`
`fsdp_config`
`fsdp_min_num_params`
`fsdp_transformer_layer_cls_to_wrap`
`full_determinism`
`gradient_accumulation_steps`
`gradient_checkpointing`
`gradient_checkpointing_kwargs`
`greater_is_better`
`group_by_length`
`half_precision_backend`
`hub_always_push`
`hub_model_id`
`hub_private_repo`
`hub_revision`
`hub_strategy`
`hub_token`
`ignore_data_skip`
`include_inputs_for_metrics`
`include_num_input_tokens_seen`
`include_tokens_per_second`
`jit_mode_eval`
`label_names`
`label_smoothing_factor`
`learning_rate`
`length_column_name`
`liger_kernel_config`
`load_best_model_at_end`
`local_process_index`	The index of the local process used.
`local_rank`
`log_level`
`log_level_replica`
`log_on_each_node`
`logging_dir`
`logging_first_step`
`logging_nan_inf_filter`
`logging_steps`
`logging_strategy`
`lr_scheduler_type`
`max_grad_norm`
`max_steps`
`metric_for_best_model`
`mp_parameters`
`n_gpu`	The number of GPUs used by this process.
`neftune_noise_alpha`
`no_cuda`
`num_train_epochs`
`optim`
`optim_args`
`optim_target_modules`
`output_dir`
`overwrite_output_dir`
`parallel_mode`	The current mode used for parallelism if multiple GPUs/TPU cores are available.
`past_index`
`per_device_eval_batch_size`
`per_device_train_batch_size`
`per_gpu_eval_batch_size`
`per_gpu_train_batch_size`
`place_model_on_device`	Can be subclassed and overridden for some specific integrations.
`prediction_loss_only`
`process_index`	The index of the current process used.
`push_to_hub`
`push_to_hub_model_id`
`push_to_hub_organization`
`push_to_hub_token`
`ray_scope`
`remove_unused_columns`
`report_to`
`restore_callback_states_from_checkpoint`
`resume_from_checkpoint`
`run_name`
`save_on_each_node`
`save_only_model`
`save_safetensors`
`save_steps`
`save_strategy`
`save_total_limit`
`seed`
`should_log`	Whether or not the current process should produce log.
`should_save`	Whether or not the current process should write to disk, e.g., to save models and checkpoints.
`skip_memory_metrics`
`tf32`
`torch_compile`
`torch_compile_backend`
`torch_compile_mode`
`torch_empty_cache_steps`
`torchdynamo`
`tpu_metrics_debug`
`tpu_num_cores`
`train_batch_size`	The actual batch size for training (may differ from per_gpu_train_batch_size in distributed training).
`use_cpu`
`use_ipex`
`use_legacy_prediction_loop`
`use_liger_kernel`
`use_mps_device`
`warmup_ratio`
`warmup_steps`
`weight_decay`
`world_size`	The number of processes used in parallel.
`lr_scheduler_kwargs`
`include_for_metrics`

Examples using `kaira.training.TrainingArguments`

Original DeepJSCC Model (Bourtsoulatze 2019) with Training

Deep Joint Source-Channel Coding (DeepJSCC) Model - Bourtsoulatze2019 Implementation

__init__(snr_min: float = 0.0, snr_max: float = 20.0, noise_variance_min: float = 0.1, noise_variance_max: float = 2.0, channel_uses: int | None = None, code_length: int | None = None, info_length: int | None = None, channel_type: str = 'awgn', push_to_hub: bool = False, hub_model_id: str | None = None, hub_token: str | None = None, hub_private: bool = False, hub_strategy: str = 'end', output_dir: str = './results', num_train_epochs: float = 10.0, per_device_train_batch_size: int = 32, per_device_eval_batch_size: int = 32, learning_rate: float = 0.0001, warmup_steps: int = 1000, logging_steps: int = 100, eval_steps: int = 500, save_steps: int = 1000, eval_strategy: str = 'steps', logging_strategy: str = 'steps', save_strategy: str = 'steps', **kwargs)[source]

Initialize TrainingArguments.

Parameters:

snr_min – Minimum SNR value for training
snr_max – Maximum SNR value for training
noise_variance_min – Minimum noise variance
noise_variance_max – Maximum noise variance
channel_uses – Number of channel uses
code_length – Length of the code
info_length – Length of information bits
channel_type – Type of channel simulation
push_to_hub – Whether to upload model to Hugging Face Hub
hub_model_id – Model ID for Hugging Face Hub (e.g., ‘username/model-name’)
hub_token – Hugging Face Hub authentication token
hub_private – Make the Hub repository private
hub_strategy – When to upload to Hub (‘end’ or ‘checkpoint’)
output_dir – Output directory for results
num_train_epochs – Number of training epochs
per_device_train_batch_size – Training batch size per device
per_device_eval_batch_size – Evaluation batch size per device
learning_rate – Learning rate
warmup_steps – Number of warmup steps
logging_steps – Log every X steps
eval_steps – Evaluate every X steps
save_steps – Save every X steps
eval_strategy – Evaluation strategy
logging_strategy – Logging strategy
save_strategy – Save strategy
**kwargs – Additional arguments passed to TrainingArguments

push_to_hub: bool = False

hub_model_id: str | None = None

hub_token: str | None = None

hub_strategy: HubStrategy | str = 'every_save'

classmethod from_hydra_config(hydra_cfg: DictConfig, **override_kwargs) → TrainingArguments[source]

Create TrainingArguments from Hydra configuration.

Parameters:

hydra_cfg – Hydra DictConfig containing training configuration
**override_kwargs – Additional arguments to override or add

Returns:

TrainingArguments instance

classmethod from_cli_args(args) → TrainingArguments[source]

Create TrainingArguments from command-line arguments.

Parameters:: args – Parsed command-line arguments (from argparse)
Returns:: TrainingArguments instance

to_dict() → Dict[str, Any][source]: Convert to dictionary representation.

to_hydra_config() → Any[source]

Convert to Hydra DictConfig.

Returns:: DictConfig representation
Raises:: ImportError – If Hydra is not available

get_snr_range() → tuple[source]: Get SNR range as tuple.

get_noise_variance_range() → tuple[source]: Get noise variance range as tuple.

accelerator_config: dict | str | None = None

adafactor: bool = False

adam_beta1: float = 0.9

adam_beta2: float = 0.999

adam_epsilon: float = 1e-08

auto_find_batch_size: bool = False

average_tokens_across_devices: bool | None = False

batch_eval_metrics: bool = False

bf16: bool = False

bf16_full_eval: bool = False

data_seed: int | None = None

dataloader_drop_last: bool = False

dataloader_num_workers: int = 0

dataloader_persistent_workers: bool = False

dataloader_pin_memory: bool = True

dataloader_prefetch_factor: int | None = None

ddp_backend: str | None = None

ddp_broadcast_buffers: bool | None = None

ddp_bucket_cap_mb: int | None = None

ddp_find_unused_parameters: bool | None = None

ddp_timeout: int = 1800

property ddp_timeout_delta: timedelta: The actual timeout for torch.distributed.init_process_group since it expects a timedelta variable.

debug: str | list[DebugOption] = ''

deepspeed: dict | str | None = None

default_optim = 'adamw_torch'

property device: device: The device used by this process.

disable_tqdm: bool | None = None

do_eval: bool = False

do_predict: bool = False

do_train: bool = False

eval_accumulation_steps: int | None = None

property eval_batch_size: int: The actual batch size for evaluation (may differ from per_gpu_eval_batch_size in distributed training).

eval_delay: float | None = 0

eval_do_concat_batches: bool = True

eval_on_start: bool = False

eval_steps: float | None = None

eval_strategy: IntervalStrategy | str = 'no'

eval_use_gather_object: bool | None = False

fp16: bool = False

fp16_backend: str = 'auto'

fp16_full_eval: bool = False

fp16_opt_level: str = 'O1'

framework = 'pt'

fsdp: list[FSDPOption] | str | None = ''

fsdp_config: dict[str, Any] | str | None = None

fsdp_min_num_params: int = 0

fsdp_transformer_layer_cls_to_wrap: str | None = None

full_determinism: bool = False

get_process_log_level()[source]

Returns the log level to be used depending on whether this process is the main process of node 0, main process of node non-0, or a non-main process.

For the main process the log level defaults to the logging level set (logging.WARNING if you didn’t do anything) unless overridden by log_level argument.

For the replica processes the log level defaults to logging.WARNING unless overridden by log_level_replica argument.

The choice between the main and replica process settings is made according to the return value of should_log.

get_warmup_steps(num_training_steps: int)[source]: Get number of steps used for a linear warmup.

gradient_accumulation_steps: int = 1

gradient_checkpointing: bool = False

gradient_checkpointing_kwargs: dict[str, Any] | str | None = None

greater_is_better: bool | None = None

group_by_length: bool = False

half_precision_backend: str = 'auto'

hub_always_push: bool = False

hub_private_repo: bool | None = None

hub_revision: str | None = None

ignore_data_skip: bool = False

include_inputs_for_metrics: bool = False

include_num_input_tokens_seen: bool | None = False

include_tokens_per_second: bool | None = False

jit_mode_eval: bool = False

label_names: list[str] | None = None

label_smoothing_factor: float = 0.0

learning_rate: float = 5e-05

length_column_name: str | None = 'length'

liger_kernel_config: dict[str, bool] | None = None

load_best_model_at_end: bool | None = False

property local_process_index: The index of the local process used.

local_rank: int = -1

log_level: str = 'passive'

log_level_replica: str = 'warning'

log_on_each_node: bool = True

logging_dir: str | None = None

logging_first_step: bool = False

logging_nan_inf_filter: bool = True

logging_steps: float = 500

logging_strategy: IntervalStrategy | str = 'steps'

lr_scheduler_type: SchedulerType | str = 'linear'

main_process_first(local=True, desc='work')[source]

A context manager for torch distributed environment where on needs to do something on the main process, while blocking replicas, and when it’s finished releasing the replicas.

One such use is for datasets’s map feature which to be efficient should be run once on the main process, which upon completion saves a cached version of results and which then automatically gets loaded by the replicas.

Parameters:

local (bool, optional, defaults to True) – if True first means process of rank 0 of each node if False first means process of rank 0 of node rank 0 In multi-node environment with a shared filesystem you most likely will want to use local=False so that only the main process of the first node will do the processing. If however, the filesystem is not shared, then the main process of each node will need to do the processing, which is the default behavior.
desc (str, optional, defaults to “work”) – a work description to be used in debug logs

max_grad_norm: float = 1.0

max_steps: int = -1

metric_for_best_model: str | None = None

mp_parameters: str = ''

property n_gpu: The number of GPUs used by this process.

Note

This will only be greater than one when you have multiple GPUs available but are not using distributed training. For distributed training, it will always be 1.

neftune_noise_alpha: float | None = None

no_cuda: bool = False

num_train_epochs: float = 3.0

optim: OptimizerNames | str = 'adamw_torch'

optim_args: str | None = None

optim_target_modules: None | str | list[str] = None

output_dir: str | None = None

overwrite_output_dir: bool = False

property parallel_mode

The current mode used for parallelism if multiple GPUs/TPU cores are available. One of:

ParallelMode.NOT_PARALLEL: no parallelism (CPU or one GPU).
ParallelMode.NOT_DISTRIBUTED: several GPUs in one single process (uses torch.nn.DataParallel).
ParallelMode.DISTRIBUTED: several GPUs, each having its own process (uses torch.nn.DistributedDataParallel).
ParallelMode.TPU: several TPU cores.

past_index: int = -1

per_device_eval_batch_size: int = 8

per_device_train_batch_size: int = 8

per_gpu_eval_batch_size: int | None = None

per_gpu_train_batch_size: int | None = None

property place_model_on_device: Can be subclassed and overridden for some specific integrations.

prediction_loss_only: bool = False

property process_index: The index of the current process used.

push_to_hub_model_id: str | None = None

push_to_hub_organization: str | None = None

push_to_hub_token: str | None = None

ray_scope: str | None = 'last'

remove_unused_columns: bool | None = True

report_to: None | str | list[str] = None

restore_callback_states_from_checkpoint: bool = False

resume_from_checkpoint: str | None = None

run_name: str | None = None

save_on_each_node: bool = False

save_only_model: bool = False

save_safetensors: bool | None = True

save_steps: float = 500

save_strategy: SaveStrategy | str = 'steps'

save_total_limit: int | None = None

seed: int = 42

property should_log: Whether or not the current process should produce log.

property should_save: Whether or not the current process should write to disk, e.g., to save models and checkpoints.

skip_memory_metrics: bool = True

tf32: bool | None = None

to_json_string()[source]: Serializes this instance to a JSON string.

to_sanitized_dict() → dict[str, Any][source]: Sanitized serialization to use with TensorBoard’s hparams

torch_compile: bool = False

torch_compile_backend: str | None = None

torch_compile_mode: str | None = None

torch_empty_cache_steps: int | None = None

torchdynamo: str | None = None

tpu_metrics_debug: bool = False

tpu_num_cores: int | None = None

property train_batch_size: int: The actual batch size for training (may differ from per_gpu_train_batch_size in distributed training).

use_cpu: bool = False

use_ipex: bool = False

use_legacy_prediction_loop: bool = False

use_liger_kernel: bool | None = False

use_mps_device: bool = False

warmup_ratio: float = 0.0

warmup_steps: int = 0

weight_decay: float = 0.0

property world_size: The number of processes used in parallel.

lr_scheduler_kwargs: dict[str, Any] | str | None

include_for_metrics: list[str]

kaira.training.TrainingArguments

Examples using kaira.training.TrainingArguments

Examples using `kaira.training.TrainingArguments`