backend

backend

Training backend abstraction and result types.

Defines TrainingBackend, the abstract base for all LLM fine-tuning backends, and NSSTrainerResult, the dataclass capturing outputs of a completed training run.

Classes:

Name	Description
NSSTrainerResult	Stores all outputs from a training run.
TrainingBackend	Abstract base class for LLM fine-tuning backends.

NSSTrainerResult(training_complete, config, log_history, adapter_path, df_train, df_ml_utility_holdout, elapsed_time) dataclass

Stores all outputs from a training run.

Attributes:

Name	Type	Description
training_complete	bool	True if training finished normally; False if stopped early.
config	SafeSynthesizerParameters	The resolved parameters used for the run.
log_history	list[dict]	Per-step log entries recorded by the HuggingFace Trainer.
adapter_path	Path	Filesystem path to the saved LoRA adapter.
df_train	DataFrame	Training DataFrame (after preprocessing).
df_ml_utility_holdout	DataFrame | None	Optional hold-out split for ML-utility evaluation.
elapsed_time	float	Wall-clock training duration in seconds.

training_complete instance-attribute

True if training finished normally; False if stopped early.

config instance-attribute

The resolved parameters used for the run.

log_history instance-attribute

Per-step log entries recorded by the HuggingFace Trainer.

adapter_path instance-attribute

Filesystem path to the saved LoRA adapter.

df_train instance-attribute

Training DataFrame (after preprocessing).

df_ml_utility_holdout instance-attribute

Optional hold-out split for ML-utility evaluation.

elapsed_time instance-attribute

Wall-clock training duration in seconds.

TrainingBackend(params, model_metadata, training_dataset=None, data_fraction=None, logging_level=logging.INFO, true_dataset_size=None, eval_dataset=None, generation_eval=False, callbacks=None, action_executor=None, workdir=None, *args, **kwargs)

Abstract base class for LLM fine-tuning backends.

Subclasses must implement every abstract method. Two concrete implementations are provided: HuggingFaceBackend (standard HuggingFace Trainer with full DP-SGD support) and UnslothTrainer (Unsloth-optimized training with lower memory usage and faster throughput, but no DP support).

Parameters:

Name	Type	Description	Default
params	SafeSynthesizerParameters	NSS pipeline configuration.	required
model_metadata	ModelMetadata	Pretrained model metadata (prompt template, sequence length, RoPE scaling, etc.).	required
training_dataset	Dataset | None	Raw tabular HuggingFace Dataset to fine-tune on. Converted to a DataFrame and run through preprocessing, schema inference, and tokenization during prepare_training_data.	None
data_fraction	float | None	Ratio of num_input_records_to_sample to total training records. Passed to the Opacus DP trainer to compute per-step sampling probability for privacy accounting.	None
logging_level	int	Python logging verbosity level.	INFO
true_dataset_size	int | None	Total number of records (or groups, when grouping is enabled) in the training set before subsampling. Used by the Opacus DP trainer alongside data_fraction for privacy accounting.	None
eval_dataset	Dataset | None	Optional raw tabular Dataset used only as a flag to enable eval-step overrides. The actual eval split passed to the Trainer is produced by the assembler (training_examples.test).	None
generation_eval	bool	If True, attach the InferenceEvalCallback during training.	False
callbacks	list[TrainerCallback] | None	Extra HuggingFace TrainerCallback instances.	None
action_executor	ActionExecutor | None	Executor for user-defined data actions (column transforms, type conversions, etc.). When provided, its preprocess phase runs on the training DataFrame before tokenization in prepare_training_data.	None
workdir	Workdir | None	Working directory for artifacts; required.	None

Raises:

Type	Description
ValueError	If workdir is not provided.

Methods:

Name	Description
prepare_training_data	Load, validate, and tokenize the training dataset.
prepare_config	Set framework-specific model-loading parameters.
prepare_params	Build training arguments and instantiate the trainer.
maybe_quantize	Apply PEFT / quantization wrapping to the loaded model.
load_model	Load the pretrained model and tokenizer.
train	Run the full training loop and populate results.
save_model	Persist the fine-tuned adapter and related artifacts.

Attributes:

Name	Type	Description
model	PreTrainedModel | PeftModel	The loaded model, with LoRA/PEFT wrapping applied after maybe_quantize.
tokenizer	PreTrainedTokenizer	Tokenizer corresponding to the pretrained model.
quant_params	dict	LoRA and optional quantization configuration populated by maybe_quantize.
load_params	dict	Raw parameters used when calling from_pretrained.
trainer_type	type[OpacusDPTrainer | Trainer | FastLanguageModel]	Trainer class to instantiate -- standard Trainer, OpacusDPTrainer for DP, or FastLanguageModel for Unsloth.
trainer	OpacusDPTrainer | Trainer	Instantiated trainer, created during prepare_params.
results	NSSTrainerResult	Outputs of the completed training run, populated by train.
training_examples	TrainingExamples	Tokenized and assembled training (and optional validation) splits.
df_train	DataFrame	Training DataFrame after preprocessing.
df_test	DataFrame | None	Hold-out DataFrame for ML-utility evaluation, or None.
training_output_dir	Path	Directory for trainer checkpoints and cache files.
train_args	TrainingArguments | dict	HuggingFace TrainingArguments (or raw dict before instantiation).
elapsed_time	float	Wall-clock training duration in seconds.
dataset_schema	dict | None	JSON schema inferred from the training DataFrame.
data_fraction	float | None	Ratio of sampled records to total, used for DP privacy accounting.
callbacks	list[TrainerCallback]	HuggingFace TrainerCallback instances attached to the trainer.
workdir	Workdir	Working directory structure for all training artifacts.

Source code in src/nemo_safe_synthesizer/training/backend.py

def __init__(
    self,
    params: SafeSynthesizerParameters,
    model_metadata: ModelMetadata,
    training_dataset: Dataset | None = None,
    data_fraction: float | None = None,
    logging_level: int = logging.INFO,
    true_dataset_size: int | None = None,
    eval_dataset: Dataset | None = None,
    generation_eval: bool = False,
    callbacks: list[TrainerCallback] | None = None,
    action_executor: ActionExecutor | None = None,
    workdir: Workdir | None = None,
    *args,
    **kwargs,
):
    self.params = params
    self.model_metadata = model_metadata
    self.dataset_schema = None
    self.framework_load_params: dict = {}
    self.data_fraction = data_fraction
    self.logging_level = logging_level
    self.true_dataset_size = true_dataset_size
    self.eval_dataset = eval_dataset
    self.training_dataset = training_dataset
    self.generation_eval = generation_eval
    self.callbacks = callbacks or []
    self.action_executor = action_executor
    if not workdir:
        raise ValueError("workdir is required")
    self.workdir = workdir

model instance-attribute

The loaded model, with LoRA/PEFT wrapping applied after maybe_quantize.

tokenizer instance-attribute

Tokenizer corresponding to the pretrained model.

quant_params instance-attribute

LoRA and optional quantization configuration populated by maybe_quantize.

load_params instance-attribute

Raw parameters used when calling from_pretrained.

trainer_type instance-attribute

Trainer class to instantiate -- standard Trainer, OpacusDPTrainer for DP, or FastLanguageModel for Unsloth.

trainer instance-attribute

Instantiated trainer, created during prepare_params.

results instance-attribute

Outputs of the completed training run, populated by train.

training_examples instance-attribute

Tokenized and assembled training (and optional validation) splits.

df_train instance-attribute

Training DataFrame after preprocessing.

df_test instance-attribute

Hold-out DataFrame for ML-utility evaluation, or None.

training_output_dir instance-attribute

Directory for trainer checkpoints and cache files.

train_args instance-attribute

HuggingFace TrainingArguments (or raw dict before instantiation).

elapsed_time instance-attribute

Wall-clock training duration in seconds.

dataset_schema = None instance-attribute

JSON schema inferred from the training DataFrame.

data_fraction = data_fraction instance-attribute

Ratio of sampled records to total, used for DP privacy accounting.

callbacks = callbacks or [] instance-attribute

HuggingFace TrainerCallback instances attached to the trainer.

workdir = workdir instance-attribute

Working directory structure for all training artifacts.

prepare_training_data() abstractmethod

Load, validate, and tokenize the training dataset.

Runs auto-config resolution, validates groupby/orderby columns, applies time-series processing and action_executor preprocessing, then assembles tokenized training examples. Populates training_examples, dataset_schema, df_train, and data_fraction.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def prepare_training_data(self):
    """Load, validate, and tokenize the training dataset.

    Runs auto-config resolution, validates groupby/orderby columns,
    applies time-series processing and ``action_executor`` preprocessing,
    then assembles tokenized training examples. Populates
    ``training_examples``, ``dataset_schema``, ``df_train``, and
    ``data_fraction``.
    """
    ...

prepare_config() abstractmethod

Set framework-specific model-loading parameters.

Builds the framework_load_params dict consumed by from_pretrained, including device map, attention implementation, quantization config, and RoPE scaling. Idempotent -- returns immediately if parameters are already prepared.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def prepare_config(self):
    """Set framework-specific model-loading parameters.

    Builds the ``framework_load_params`` dict consumed by
    ``from_pretrained``, including device map, attention implementation,
    quantization config, and RoPE scaling. Idempotent -- returns
    immediately if parameters are already prepared.
    """
    ...

prepare_params() abstractmethod

Build training arguments and instantiate the trainer.

Constructs TrainingArguments from the resolved config, configures DP or standard training, selects the data collator, and creates the trainer instance with attached callbacks.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def prepare_params(self):
    """Build training arguments and instantiate the trainer.

    Constructs ``TrainingArguments`` from the resolved config, configures
    DP or standard training, selects the data collator, and creates the
    ``trainer`` instance with attached callbacks.
    """
    ...

maybe_quantize(**quant_params) abstractmethod

Apply PEFT / quantization wrapping to the loaded model.

Configures LoRA parameters (rank, alpha, target modules) and wraps model as a PEFT model. When quantization is enabled, applies k-bit preparation before wrapping.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def maybe_quantize(self, **quant_params: dict):
    """Apply PEFT / quantization wrapping to the loaded model.

    Configures LoRA parameters (rank, alpha, target modules) and wraps
    ``model`` as a PEFT model. When quantization is enabled, applies
    k-bit preparation before wrapping.
    """
    ...

load_model() abstractmethod

Load the pretrained model and tokenizer.

Calls prepare_config to resolve loading parameters, loads the model and tokenizer via the framework-specific loader (e.g. AutoModelForCausalLM or FastLanguageModel), then applies PEFT/quantization via maybe_quantize.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def load_model(self):
    """Load the pretrained model and tokenizer.

    Calls ``prepare_config`` to resolve loading parameters, loads the
    model and tokenizer via the framework-specific loader (e.g.
    ``AutoModelForCausalLM`` or ``FastLanguageModel``), then applies
    PEFT/quantization via ``maybe_quantize``.
    """
    ...

train() abstractmethod

Run the full training loop and populate results.

Orchestrates the end-to-end pipeline: prepares data, builds training arguments, runs the trainer, and saves artifacts. Populates results with the training outcome.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def train(self):
    """Run the full training loop and populate ``results``.

    Orchestrates the end-to-end pipeline: prepares data, builds
    training arguments, runs the trainer, and saves artifacts.
    Populates ``results`` with the training outcome.
    """
    ...

save_model() abstractmethod

Persist the fine-tuned adapter and related artifacts.

Saves the LoRA adapter weights, model metadata, dataset schema, and resolved config to the workdir. Optionally frees the model from GPU memory after saving.

Source code in src/nemo_safe_synthesizer/training/backend.py

@abc.abstractmethod
def save_model(self):
    """Persist the fine-tuned adapter and related artifacts.

    Saves the LoRA adapter weights, model metadata, dataset schema,
    and resolved config to the ``workdir``. Optionally frees the
    model from GPU memory after saving.
    """
    ...