mostlyai-engine 2.6.1__tar.gz → 2.6.2__tar.gz

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/PKG-INFO

@@ -1,11 +1,11 @@
 Metadata-Version: 2.4
 Name: mostlyai-engine
-Version: 2.6.1
+Version: 2.6.2
 Summary: Synthetic Data Engine
 Project-URL: homepage, https://github.com/mostly-ai/mostlyai-engine
 Project-URL: repository, https://github.com/mostly-ai/mostlyai-engine
+Project-URL: issues, https://github.com/mostly-ai/mostlyai-engine/issues
 Project-URL: documentation, https://mostly-ai.github.io/mostlyai-engine/
-Author-email: MOSTLY AI <dev@mostly.ai>
 License-Expression: Apache-2.0
 License-File: LICENSE
 Classifier: Development Status :: 5 - Production/Stable
@@ -29,7 +29,7 @@ Requires-Dist: huggingface-hub[hf-xet]>=0.30.2
 Requires-Dist: joblib>=1.4.2
 Requires-Dist: json-repair>=0.47.0
 Requires-Dist: numpy>=2.0.0
-Requires-Dist: opacus>=1.5.4
+Requires-Dist: opacus>=1.6.0
 Requires-Dist: pandas>=2.2.0
 Requires-Dist: peft>=0.18.2
 Requires-Dist: psutil<6,>=5.9.5

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/mostlyai/engine/__init__.py

@@ -34,7 +34,7 @@ __all__ = [
     "TabularARGN",
     "LanguageModel",
 ]
-__version__ = "2.6.1"
+__version__ = "2.6.2"
 
 # suppress specific warning related to os.fork() in multi-threaded processes
 warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*multi-threaded.*fork.*")

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/mostlyai/engine/_language/training.py

@@ -31,7 +31,7 @@ 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.grad_sample import GradSampleHooks, register_grad_sampler
 from opacus.utils.batch_memory_manager import wrap_data_loader
 from peft import LoraConfig, PeftModel
 from torch import nn
@@ -626,6 +626,7 @@ def train(
             # this can help accelerate GPU compute
             torch.backends.cudnn.benchmark = True
 
+        dp_grad_sample_hooks: GradSampleHooks | None = None
         if with_dp:
             if isinstance(differential_privacy, DifferentialPrivacyConfig):
                 dp_config = differential_privacy.model_dump()
@@ -650,18 +651,21 @@ def train(
                 privacy_engine.accountant.load_state_dict(
                     torch.load(workspace.model_dp_accountant_path, map_location=device, weights_only=True),
                 )
-            # Opacus will return the modified objects
-            # - model: wrapped in GradSampleModule and contains additional hooks for computing per-sample gradients
-            # - optimizer: wrapped in DPOptimizer and will do different operations during virtual steps and logical steps
-            # - dataloader: the dataloader with batch_sampler=UniformWithReplacementSampler (for Poisson sampling)
-            model, optimizer, trn_dataloader = privacy_engine.make_private(
+            # Opacus returns GradSampleHooks when wrap_model=False: hooks attach to the original module so HF /
+            # Transformers sees an unwrapped PreTrainedModel (requires Opacus >= 1.6).
+            # - dp_grad_sample_hooks: must call .cleanup() after training to remove backward hooks and param attrs
+            # - optimizer: wrapped in DPOptimizer (virtual vs logical steps)
+            # - dataloader: UniformWithReplacementSampler when poisson_sampling=True
+            dp_grad_sample_hooks, optimizer, trn_dataloader = privacy_engine.make_private(
                 module=model,
                 optimizer=optimizer,
                 data_loader=trn_dataloader,
                 noise_multiplier=dp_config.get("noise_multiplier"),
                 max_grad_norm=dp_config.get("max_grad_norm"),
                 poisson_sampling=True,
+                wrap_model=False,
             )
+            model = dp_grad_sample_hooks._module
             # this further wraps the dataloader with batch_sampler=BatchSplittingSampler to achieve gradient accumulation
             # it will split the sampled logical batches into smaller sub-batches with batch_size
             trn_dataloader = wrap_data_loader(
@@ -835,6 +839,9 @@ def train(
             if total_training_time > max_training_time:
                 do_stop = True
 
+        if dp_grad_sample_hooks is not None:
+            dp_grad_sample_hooks.cleanup()
+
         # no checkpoint is saved yet because the training stopped before the first epoch ended
         if not model_checkpoint.has_saved_once():
             _LOG.info("saving model weights, as none were saved so far")

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/mostlyai/engine/_tabular/argn.py

@@ -289,7 +289,8 @@ class SequentialContextEmbedders(Embedders):
         mask = None
         for sub_col in self.cardinalities:
             xs = torch.as_tensor(x[sub_col], device=self.device)
-            xs = torch.nested.to_padded_tensor(xs, padding=-1)
+            if xs.is_nested:
+                xs = torch.nested.to_padded_tensor(xs, padding=-1)
             mask = (xs != -1).squeeze(-1)
             xs = torch.where(xs == -1, torch.tensor(0), xs)
             xs = self.get(sub_col)(xs)

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/mostlyai/engine/_tabular/training.py

@@ -138,10 +138,20 @@ class BatchCollator:
     For sequence data, it will sample subsequences with lengths up to max_sequence_window.
     """
 
-    def __init__(self, is_sequential: bool, max_sequence_window: int | None, device: torch.device):
+    def __init__(
+        self,
+        is_sequential: bool,
+        max_sequence_window: int | None,
+        device: torch.device,
+        *,
+        use_nested_ctxseq: bool = True,
+    ):
         self.is_sequential = is_sequential
         self.max_sequence_window = max_sequence_window
         self.device = device
+        # Opacus per-sample gradients do not support NestedTensor on CPU/CUDA; use padded
+        # dense tensors for CTXSEQ when training with DP (see test_tabular_sequential DP path).
+        self.use_nested_ctxseq = use_nested_ctxseq
 
     def __call__(self, batch: list[dict]) -> dict[str, torch.Tensor]:
         batch = pd.DataFrame(batch)
@@ -177,15 +187,26 @@ class BatchCollator:
                     dim=-1,
                 )
             elif column.startswith(CTXSEQ):
-                # construct row tensors and convert the list to nested column tensor
-                tensors[column] = torch.unsqueeze(
-                    torch.nested.as_nested_tensor(
-                        [torch.tensor(row, dtype=torch.int64, device=self.device) for row in batch[column]],
-                        dtype=torch.int64,
-                        device=self.device,
-                    ),
-                    dim=-1,
-                )
+                if self.use_nested_ctxseq:
+                    # construct row tensors and convert the list to nested column tensor
+                    tensors[column] = torch.unsqueeze(
+                        torch.nested.as_nested_tensor(
+                            [torch.tensor(row, dtype=torch.int64, device=self.device) for row in batch[column]],
+                            dtype=torch.int64,
+                            device=self.device,
+                        ),
+                        dim=-1,
+                    )
+                else:
+                    # padded batch (variable-length rows); -1 marks padding (matches SequentialContextEmbedders)
+                    tensors[column] = torch.unsqueeze(
+                        torch.tensor(
+                            np.array(list(zip_longest(*batch[column], fillvalue=-1))).T,
+                            dtype=torch.int64,
+                            device=self.device,
+                        ),
+                        dim=-1,
+                    )
         return tensors
 
     @staticmethod
@@ -544,7 +565,10 @@ def train(
 
         # and see if it's possible to make it compatible with DP
         batch_collator = BatchCollator(
-            is_sequential=is_sequential, max_sequence_window=max_sequence_window, device=device
+            is_sequential=is_sequential,
+            max_sequence_window=max_sequence_window,
+            device=device,
+            use_nested_ctxseq=not with_dp,
         )
         disable_progress_bar()
         trn_dataset = load_dataset("parquet", data_files=[str(p) for p in workspace.encoded_data_trn.fetch_all()])[

{mostlyai_engine-2.6.1 → mostlyai_engine-2.6.2}/pyproject.toml

@@ -1,8 +1,7 @@
 [project]
 name = "mostlyai-engine"
-version = "2.6.1"
+version = "2.6.2"
 description = "Synthetic Data Engine"
-authors = [{ name = "MOSTLY AI", email = "dev@mostly.ai" }]
 requires-python = ">=3.11,<3.14"
 readme = "README.md"
 license = "Apache-2.0"
@@ -37,7 +36,7 @@ dependencies = [
     "accelerate>=1.5.0",
     "peft>=0.18.2",  # transformers 5.7+ checks min PEFT in model.add_adapter (integrations/peft.py)
     "huggingface-hub[hf-xet]>=0.30.2",
-    "opacus>=1.5.4",
+    "opacus>=1.6.0",
     "xgrammar>=0.1.32,<1.0.0",  # aligned with vllm 0.20
     "json-repair>=0.47.0",
     "torch>=2.11.0,<2.12.0",
@@ -75,6 +74,7 @@ docs = [
 [project.urls]
 homepage = "https://github.com/mostly-ai/mostlyai-engine"
 repository = "https://github.com/mostly-ai/mostlyai-engine"
+issues = "https://github.com/mostly-ai/mostlyai-engine/issues"
 documentation = "https://mostly-ai.github.io/mostlyai-engine/"
 
 [tool.uv]