mostlyai-engine 1.4.5__tar.gz → 1.4.7__tar.gz

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 1.4.5
+Version: 1.4.7
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
@@ -29,23 +29,23 @@ Requires-Dist: datasets>=3.0.0
 Requires-Dist: huggingface-hub[hf-xet]>=0.30.2
 Requires-Dist: joblib>=1.4.2
 Requires-Dist: json-repair<0.47.0,>=0.30.0
-Requires-Dist: numpy>=1.26.3
-Requires-Dist: opacus>=1.5.2
+Requires-Dist: numpy>=2.0.0
+Requires-Dist: opacus>=1.5.4
 Requires-Dist: pandas~=2.2.0
 Requires-Dist: peft>=0.12.0
 Requires-Dist: psutil<6,>=5.9.5
 Requires-Dist: pyarrow>=16.0.0
 Requires-Dist: setuptools>=77.0.3
 Requires-Dist: tokenizers>=0.21.0
-Requires-Dist: torch<2.6.1,>=2.6.0
-Requires-Dist: torchaudio<2.6.1,>=2.6.0
-Requires-Dist: torchvision<0.21.1,>=0.21.0
+Requires-Dist: torch<2.7.1,>=2.7.0
+Requires-Dist: torchaudio<2.7.1,>=2.7.0
+Requires-Dist: torchvision<0.22.1,>=0.22.0
 Requires-Dist: transformers>=4.51.0
-Requires-Dist: xgrammar>=0.1.18
+Requires-Dist: xgrammar>=0.1.19
 Provides-Extra: gpu
 Requires-Dist: bitsandbytes==0.42.0; (sys_platform == 'darwin') and extra == 'gpu'
 Requires-Dist: bitsandbytes>=0.45.5; (sys_platform == 'linux') and extra == 'gpu'
-Requires-Dist: vllm==0.8.5.post1; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
+Requires-Dist: vllm==0.9.1; (sys_platform == 'linux' or sys_platform == 'darwin') and extra == 'gpu'
 Description-Content-Type: text/markdown
 
 # Synthetic Data Engine 💎
@@ -91,7 +91,7 @@ pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-pip install -U torch==2.6.0+cpu torchvision==0.21.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+pip install -U torch==2.7.0+cpu torchvision==0.22.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## Quick start

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/README.md

@@ -41,7 +41,7 @@ pip install -U 'mostlyai-engine[gpu]'
 On Linux, one can explicitly install the CPU-only variant of torch together with `mostlyai-engine`:
 
 ```bash
-pip install -U torch==2.6.0+cpu torchvision==0.21.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
+pip install -U torch==2.7.0+cpu torchvision==0.22.0+cpu mostlyai-engine --extra-index-url https://download.pytorch.org/whl/cpu
 ```
 
 ## Quick start

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/mostlyai/engine/__init__.py

@@ -22,7 +22,7 @@ from mostlyai.engine.splitting import split
 from mostlyai.engine.training import train
 
 __all__ = ["split", "analyze", "encode", "train", "generate", "init_logging", "set_random_state"]
-__version__ = "1.4.5"
+__version__ = "1.4.7"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/mostlyai/engine/_language/common.py

@@ -18,7 +18,15 @@ from pathlib import Path
 
 import torch
 from peft import PeftConfig, prepare_model_for_kbit_training
-from transformers import AutoConfig, AutoModelForCausalLM, BitsAndBytesConfig, PretrainedConfig, PreTrainedModel
+from transformers import (
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    BitsAndBytesConfig,
+    PretrainedConfig,
+    PreTrainedModel,
+)
 
 from mostlyai.engine._language.lstm import LSTMFromScratchConfig
 
@@ -35,7 +43,7 @@ def is_bf16_supported(device: torch.device) -> bool:
 
 
 def get_attention_implementation(config: PretrainedConfig) -> str | None:
-    model_cls = AutoModelForCausalLM._model_mapping[type(config)]
+    model_cls = AutoModel._model_mapping[type(config)]
     attn_implementation = None
     if getattr(model_cls, "_supports_sdpa", False):
         attn_implementation = "sdpa"
@@ -45,6 +53,7 @@ def get_attention_implementation(config: PretrainedConfig) -> str | None:
 def load_base_model_and_config(
     model_id_or_path: str | Path, device: torch.device, is_peft_adapter: bool, is_training: bool
 ) -> tuple[PreTrainedModel, PretrainedConfig]:
+    # opacus DP does not support parallel/sharded training
     model_id_or_path = str(model_id_or_path)
     if is_peft_adapter:
         # get the base model name from adapter_config.json
@@ -84,13 +93,29 @@ def load_base_model_and_config(
         )
     else:
         quantization_config = None
-    model = AutoModelForCausalLM.from_pretrained(
+
+    if device.type == "cuda" and device.index is not None:
+        device_map = str(device)
+    else:
+        device_map = "auto"
+
+    if hasattr(config, "text_config") and hasattr(config, "vision_config"):
+        config.text_config.use_cache = use_cache
+        config.text_config.attn_implementation = attn_implementation
+        auto_model_cls = AutoModelForImageTextToText
+    elif hasattr(config, "use_cache"):
+        config.use_cache = use_cache
+        config.attn_implementation = attn_implementation
+        auto_model_cls = AutoModelForCausalLM
+    else:
+        raise ValueError("Unsupported model")
+
+    model = auto_model_cls.from_pretrained(
         model_id_or_path,
-        torch_dtype=torch_dtype,
-        attn_implementation=attn_implementation,
-        use_cache=use_cache,
-        device_map=device,
+        config=config,
+        device_map=device_map,
         quantization_config=quantization_config,
+        torch_dtype=torch_dtype,
     )
     if quantization_config:
         # convert all non-kbit layers to float32

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/mostlyai/engine/_language/engine/vllm_engine.py

@@ -152,6 +152,7 @@ class VLLMEngine(LanguageEngine):
             # enforce_eager=True,  # results in big slowdown, but is needed when running pytest locally
             swap_space=0,
             disable_log_stats=True,
+            tensor_parallel_size=torch.cuda.device_count(),
         )
         self.tokenizer = AutoTokenizer.from_pretrained(
             model_path,

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/mostlyai/engine/_language/training.py

@@ -28,13 +28,16 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 import torch
+from accelerate import Accelerator, FullyShardedDataParallelPlugin
 from datasets import Dataset, DatasetDict, disable_progress_bar, load_dataset
+from huggingface_hub import get_safetensors_metadata
 from opacus import GradSampleModule, PrivacyEngine
 from opacus.accountants import GaussianAccountant, PRVAccountant, RDPAccountant
 from opacus.grad_sample import register_grad_sampler
 from opacus.utils.batch_memory_manager import wrap_data_loader
 from peft import LoraConfig, PeftModel
 from torch import nn
+from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
 from torch.nn import CrossEntropyLoss
 from torch.optim.lr_scheduler import LRScheduler
 from torch.utils.data import DataLoader
@@ -212,6 +215,12 @@ def _calculate_val_loss(model: PreTrainedModel | GradSampleModule, val_dataloade
     return val_loss_avg.item()
 
 
+def get_num_model_params(model: str) -> int:
+    metadata = get_safetensors_metadata(model)
+    no_of_model_params = next(iter(metadata.parameter_count.values()))
+    return no_of_model_params
+
+
 def _calculate_max_tokens(tokenized_trn_dataset: Dataset) -> int:
     max_tokens = 0
     for example in tokenized_trn_dataset:
@@ -294,13 +303,37 @@ def train(
     ) as progress:
         _LOG.info(f"numpy={version('numpy')}, pandas={version('pandas')}")
         _LOG.info(f"torch={version('torch')}, opacus={version('opacus')}")
-        _LOG.info(f"transformers={version('transformers')}, peft={version('peft')}")
+        _LOG.info(f"transformers={version('transformers')}, accelerate={version('accelerate')}, peft={version('peft')}")
+        with_dp = differential_privacy is not None
         device = (
             torch.device(device)
             if device is not None
             else (torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu"))
         )
+
+        single_gpu_threshold = 7_000_000_000
+        if (
+            device.type == "cuda"
+            and device.index is None
+            and (
+                with_dp or model == LSTMFromScratchConfig.model_id or get_num_model_params(model) < single_gpu_threshold
+            )
+        ):
+            device = torch.device("cuda:0")
+            _LOG.info("device set to single gpu (cuda:0) because model is too small or differential privacy is enabled")
+
+        if not with_dp:
+            if device.type == "cuda":
+                fsdp_plugin = FullyShardedDataParallelPlugin(
+                    state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=False),
+                    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False),
+                )
+            else:
+                fsdp_plugin = None
+            accelerator = Accelerator(fsdp_plugin=fsdp_plugin, cpu=device.type == "cpu")
+
         _LOG.info(f"{device=}")
+        _LOG.info(f"{torch.cuda.device_count()=}")
         bf16_supported = is_bf16_supported(device)
         _LOG.info(f"{bf16_supported=}")
         use_mixed_precision = bf16_supported and model != LSTMFromScratchConfig.model_id
@@ -324,7 +357,7 @@ def train(
             max_epochs = max_epochs_cap
         else:
             _LOG.info(f"{max_epochs=}")
-        with_dp = differential_privacy is not None
+
         _LOG.info(f"{with_dp=}")
         _LOG.info(f"{model_state_strategy=}")
 
@@ -566,6 +599,11 @@ def train(
             min_lr=0.1 * initial_lr,
             # threshold=0,  # if we prefer to completely mimic the behavior of previous implementation
         )
+        is_reduce_lr_on_plateau = isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau)
+
+        if not with_dp:
+            model, optimizer, lr_scheduler = accelerator.prepare(model, optimizer, lr_scheduler)
+
         if (
             model_state_strategy == ModelStateStrategy.resume
             and model_checkpoint.optimizer_and_lr_scheduler_paths_exist()
@@ -626,6 +664,7 @@ def train(
         else:
             privacy_engine = None
             dp_config, dp_total_delta, dp_accountant = None, None, None
+            trn_dataloader = accelerator.prepare(trn_dataloader)
 
         progress_message = None
         start_trn_time = time.time()
@@ -667,8 +706,12 @@ def train(
                         outputs = model(**step_data)
                     # FIXME approximation, should be divided by total sum of number of tokens in the batch
                     #  as in _calculate_per_label_losses, also the final sample may be smaller than the batch size.
-                    step_loss = outputs.loss / (1 if with_dp else gradient_accumulation_steps)
-                    step_loss.backward()
+                    if with_dp:
+                        step_loss = outputs.loss
+                        step_loss.backward()
+                    else:
+                        step_loss = outputs.loss / gradient_accumulation_steps
+                        accelerator.backward(step_loss)
                 accumulated_steps += 1
                 # explicitly count the number of processed samples as the actual batch size can vary when DP is on
                 samples += step_data["input_ids"].shape[0]
@@ -686,8 +729,9 @@ def train(
             current_lr = optimizer.param_groups[0][
                 "lr"
             ]  # currently assume that we have the same lr for all param groups
+
             # only the scheduling for ReduceLROnPlateau is postponed until the metric becomes available
-            if not isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+            if not is_reduce_lr_on_plateau:
                 lr_scheduler.step()
 
             # do validation
@@ -727,7 +771,7 @@ def train(
                 # check for early stopping
                 do_stop = early_stopper(val_loss=val_loss) or has_exceeded_dp_max_epsilon
                 # scheduling for ReduceLROnPlateau
-                if isinstance(lr_scheduler, torch.optim.lr_scheduler.ReduceLROnPlateau):
+                if is_reduce_lr_on_plateau:
                     lr_scheduler.step(metrics=val_loss)
 
             # log progress, either by time or by steps, whatever is shorter
@@ -788,7 +832,7 @@ def train(
         if not model_checkpoint.has_saved_once():
             _LOG.info("saving model weights, as none were saved so far")
             model_checkpoint.save_checkpoint(
-                model=model,
+                model=model if with_dp else accelerator.unwrap_model(model),
                 optimizer=optimizer,
                 lr_scheduler=lr_scheduler,
                 dp_accountant=privacy_engine.accountant if with_dp else None,

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/mostlyai/engine/_tabular/generation.py

@@ -758,15 +758,7 @@ def generate(
         _LOG.info(f"{gen_column_order=}")
         if not enable_flexible_generation:
             # check if resolved column order is the same as the one from training
-            trn_column_order = [SLEN_SIDX_SDEC_COLUMN] if is_sequential else []
-            trn_column_order += [
-                get_argn_name(
-                    argn_processor=tgt_stats["columns"][col][ARGN_PROCESSOR],
-                    argn_table=tgt_stats["columns"][col][ARGN_TABLE],
-                    argn_column=tgt_stats["columns"][col][ARGN_COLUMN],
-                )
-                for col in tgt_stats["columns"].keys()
-            ]
+            trn_column_order = get_columns_from_cardinalities(tgt_cardinalities)
             _LOG.info(f"{trn_column_order=}")
             if gen_column_order != trn_column_order:
                 raise ValueError(

{mostlyai_engine-1.4.5 → mostlyai_engine-1.4.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mostlyai-engine"
-version = "1.4.5"
+version = "1.4.7"
 description = "Synthetic Data Engine"
 authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.10"
@@ -25,8 +25,8 @@ classifiers = [
 ]
 
 dependencies = [
-    "setuptools>=77.0.3",  # similar to vllm 0.8.5.post1
-    "numpy>=1.26.3",
+    "setuptools>=77.0.3",  # similar to vllm 0.9.1
+    "numpy>=2.0.0",
     "pandas~=2.2.0",
     "pyarrow>=16.0.0",
     "joblib>=1.4.2",
@@ -37,19 +37,19 @@ dependencies = [
     "accelerate>=1.5.0",
     "peft>=0.12.0",
     "huggingface-hub[hf-xet]>=0.30.2",
-    "opacus>=1.5.2",  # switch to 1.5.4 (once released) to allow numpy 2
-    "xgrammar>=0.1.18",  # for vllm 0.8.5.post1 compatibility
+    "opacus>=1.5.4",
+    "xgrammar>=0.1.19",  # for vllm 0.9.1 compatibility
     "json-repair>=0.30.0, <0.47.0",  # fixes errors in tests when building engine 1.4.4
-    "torch>=2.6.0,<2.6.1",
-    "torchaudio>=2.6.0,<2.6.1",  # for vllm 0.8.5.post1 compatibility
-    "torchvision>=0.21.0,<0.21.1"  # for vllm 0.8.5.post1 compatibility
+    "torch>=2.7.0,<2.7.1",  # for vllm 0.9.1 compatibility
+    "torchaudio>=2.7.0,<2.7.1",  # for vllm 0.9.1 compatibility
+    "torchvision>=0.22.0,<0.22.1"  # for vllm 0.9.1 compatibility
 ]
 
 [project.optional-dependencies]
 gpu = [
     "bitsandbytes==0.42.0; sys_platform == 'darwin'",
     "bitsandbytes>=0.45.5; sys_platform == 'linux'",
-    "vllm==0.8.5.post1; sys_platform == 'linux' or sys_platform == 'darwin'",
+    "vllm==0.9.1; sys_platform == 'linux' or sys_platform == 'darwin'",
 ]
 
 [dependency-groups]