From c070f71e16bc2e32f219e83886d59eff4f52dfae Mon Sep 17 00:00:00 2001 From: petmycat <160725257+petmycat@users.noreply.github.com> Date: Fri, 14 Aug 2026 20:57:07 +0800 Subject: [PATCH] Fix/oo mfix1 (#8) * fix: OOM error 1 * fix: OOMfix2 --------- Co-authored-by: James <1561862923@qq.com> --- extensions_built_in/sd_trainer/SDTrainer.py | 40 ++++++++++++++----- .../test_sd_trainer_three_phase_runtime.py | 26 +++++++++++- testing/test_trigger_binding_artifacts.py | 29 ++++++++++++++ toolkit/trigger_binding_artifacts.py | 10 ++++- 4 files changed, 90 insertions(+), 15 deletions(-) diff --git a/extensions_built_in/sd_trainer/SDTrainer.py b/extensions_built_in/sd_trainer/SDTrainer.py index 5144ecd4..b49a1014 100644 --- a/extensions_built_in/sd_trainer/SDTrainer.py +++ b/extensions_built_in/sd_trainer/SDTrainer.py @@ -119,6 +119,7 @@ class SDTrainer(BaseSDTrainProcess): self._trigger_binding_initial_parameters = {} self._trigger_binding_prompt_encoder = None self._trigger_binding_last_metrics = {} + self._trigger_binding_last_metrics_written_step = None self.dfe: Optional[DiffusionFeatureExtractor] = None self.unconditional_embeds = None @@ -1757,6 +1758,33 @@ class SDTrainer(BaseSDTrainProcess): setter(previous) return runtime_context() + def _write_trigger_binding_metrics(self, loss=None): + if not self.three_phase_enabled or self.runtime_phase not in {'a1', 'a2'}: + return + if self._trigger_binding_last_metrics_written_step == self.step_num: + return + artifact_config = getattr( + self.three_phase_trigger_training.artifacts, + f'phase_{self.runtime_phase}', + ) + phase_root = os.path.join( + self.three_phase_trigger_training.run_root or self.save_root, + f'phase_{self.runtime_phase}', + ) + metrics_path = os.path.join(phase_root, artifact_config.metrics_file) + os.makedirs(os.path.dirname(metrics_path), exist_ok=True) + record = { + 'phase': self.runtime_phase, + 'step': self.step_num, + 'metrics': self._trigger_binding_last_metrics, + } + if loss is not None: + record['loss'] = float(loss.detach().item() if torch.is_tensor(loss) else loss) + import json + with open(metrics_path, 'a', encoding='utf-8') as handle: + handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + '\n') + self._trigger_binding_last_metrics_written_step = self.step_num + def _calculate_trigger_binding_loss( self, noisy_latents, @@ -1875,6 +1903,7 @@ class SDTrainer(BaseSDTrainProcess): if isinstance(value, (int, float)): self.additional_logs[f'phase/{self.runtime_phase}/{key}'] = float(value) self.additional_logs[f'phase/{self.runtime_phase}/loss'] = float(loss.detach().item()) + self._write_trigger_binding_metrics(loss) return loss def _encode_tst_prompt_variants(self, batch, trigger_prompts, decoy_prompts, dtype): @@ -3131,17 +3160,6 @@ class SDTrainer(BaseSDTrainProcess): 'parameter_change_proof': proof, }, ) - metrics_path = os.path.join(phase_root, artifact_config.metrics_file) - os.makedirs(os.path.dirname(metrics_path), exist_ok=True) - with open(metrics_path, 'a', encoding='utf-8') as handle: - import json - handle.write(json.dumps({ - 'phase': self.runtime_phase, - 'step': self.step_num, - 'metrics': self._trigger_binding_last_metrics, - 'parameter_change_proof': proof, - }, sort_keys=True) + '\n') - def hook_train_loop(self, batch: Union[DataLoaderBatchDTO, List[DataLoaderBatchDTO]]): if isinstance(batch, list): batch_list = batch diff --git a/testing/test_sd_trainer_three_phase_runtime.py b/testing/test_sd_trainer_three_phase_runtime.py index 58c2b09d..edc2724d 100644 --- a/testing/test_sd_trainer_three_phase_runtime.py +++ b/testing/test_sd_trainer_three_phase_runtime.py @@ -3,6 +3,7 @@ import contextlib import importlib import inspect import os +import tempfile import types import unittest from pathlib import Path @@ -19,8 +20,8 @@ def _load_runtime_methods(): names = { 'three_phase_enabled', '_load_trigger_binding_modules', '_call_supported', '_first_callable', '_phase_config', '_activator_component_flags', '_configure_phase_trainability', - 'hook_add_extra_train_params', '_activator_mode', '_calculate_trigger_binding_loss', - '_install_trigger_binding_prompt_encoder', 'encode_static_prompt', + 'hook_add_extra_train_params', '_activator_mode', '_write_trigger_binding_metrics', + '_calculate_trigger_binding_loss', '_install_trigger_binding_prompt_encoder', 'encode_static_prompt', } selected = [node for node in class_node.body if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name in names] module = ast.Module(body=[ast.ClassDef(name='SDTrainerRuntimeHarness', bases=[], keywords=[], body=selected, decorator_list=[])], type_ignores=[]) @@ -196,6 +197,26 @@ class ThreePhaseRuntimeTest(unittest.TestCase): with self.assertRaisesRegex(ValueError, 'every training caption must contain'): trainer.sd.get_prompt_embeds(['caption without the required token']) + def test_phase_metrics_are_written_independently_and_once_per_step(self): + trainer = self._trainer('a1') + trainer.step_num = 7 + trainer._trigger_binding_last_metrics = {'gain': 0.25} + trainer._trigger_binding_last_metrics_written_step = None + with tempfile.TemporaryDirectory() as temp_dir: + trainer.save_root = temp_dir + trainer.three_phase_trigger_training.run_root = temp_dir + trainer.three_phase_trigger_training.artifacts = SimpleNamespace( + phase_a1=SimpleNamespace(metrics_file='metrics.jsonl'), + ) + trainer._write_trigger_binding_metrics(torch.tensor(0.5)) + trainer._write_trigger_binding_metrics(torch.tensor(0.75)) + metrics_path = Path(temp_dir) / 'phase_a1' / 'metrics.jsonl' + records = metrics_path.read_text(encoding='utf-8').splitlines() + self.assertEqual(len(records), 1) + self.assertIn('"step": 7', records[0]) + self.assertIn('"loss": 0.5', records[0]) + self.assertIn('"gain": 0.25', records[0]) + def test_a_phase_loss_receives_shared_latent_noise_timestep_and_target(self): trainer = self._trainer('a1') trainer.device_torch = torch.device('cpu') @@ -223,6 +244,7 @@ class ThreePhaseRuntimeTest(unittest.TestCase): trainer._trigger_binding_modules = { 'losses': types.SimpleNamespace(calculate_trigger_binding_losses=fake_losses) } + trainer._write_trigger_binding_metrics = lambda _loss: None loss = trainer._calculate_trigger_binding_loss( noisy, noise, timesteps, batch, {}, 1.0, torch.float32 ) diff --git a/testing/test_trigger_binding_artifacts.py b/testing/test_trigger_binding_artifacts.py index 88403da0..d9ddb1ef 100644 --- a/testing/test_trigger_binding_artifacts.py +++ b/testing/test_trigger_binding_artifacts.py @@ -24,6 +24,7 @@ from toolkit.trigger_binding_artifacts import ( save_artifact, save_checkpoint_manifest, sha256_bytes, + tensor_sha256, source_fingerprint, ) @@ -73,6 +74,34 @@ class TriggerBindingArtifactsTest(unittest.TestCase): 64, ) + def test_tensor_hash_supports_bfloat16_scalars_empty_and_noncontiguous_tensors(self): + scalar = torch.tensor(1.25, dtype=torch.bfloat16) + vector = scalar.reshape(1) + self.assertEqual(tensor_sha256(scalar), tensor_sha256(vector)) + + empty = torch.empty(0, dtype=torch.bfloat16) + self.assertEqual(tensor_sha256(empty), sha256_bytes(b"")) + + base = torch.arange(12, dtype=torch.float32).reshape(3, 4) + noncontiguous = base.transpose(0, 1) + self.assertFalse(noncontiguous.is_contiguous()) + self.assertEqual( + tensor_sha256(noncontiguous), + tensor_sha256(noncontiguous.contiguous()), + ) + + def test_artifact_round_trip_supports_bfloat16_scalar_tensor(self): + tensors = { + "adapter.scale": torch.tensor(1.0, dtype=torch.bfloat16), + "adapter.weight": torch.ones(2, 2, dtype=torch.bfloat16), + } + with tempfile.TemporaryDirectory() as temp_dir: + path, manifest = self._save(temp_dir, tensors=tensors) + loaded, loaded_manifest = load_artifact(path, expected_type="te_adapter") + self.assertEqual(manifest, loaded_manifest) + self.assertEqual(loaded["adapter.scale"].shape, torch.Size([])) + self.assertTrue(torch.equal(loaded["adapter.scale"], tensors["adapter.scale"])) + def test_fingerprints_are_canonical_and_order_independent(self): self.assertEqual(config_fingerprint({"a": 1, "b": 2}), config_fingerprint({"b": 2, "a": 1})) self.assertNotEqual(config_fingerprint({"a": 1}), config_fingerprint({"a": 2})) diff --git a/toolkit/trigger_binding_artifacts.py b/toolkit/trigger_binding_artifacts.py index bd116438..903ef37e 100644 --- a/toolkit/trigger_binding_artifacts.py +++ b/toolkit/trigger_binding_artifacts.py @@ -129,7 +129,7 @@ def to_json_compatible(value: Any) -> Any: return {"__type__": "torch.dtype", "value": str(value).removeprefix("torch.")} if isinstance(value, torch.Tensor): tensor = value.detach().cpu().contiguous() - raw = tensor.view(torch.uint8).numpy().tobytes() + raw = _tensor_bytes(tensor) return { "__type__": "torch.Tensor", "dtype": str(tensor.dtype).removeprefix("torch."), @@ -207,7 +207,13 @@ def decode_rng_state(encoded: Any) -> Any: def _tensor_bytes(tensor: torch.Tensor) -> bytes: - return tensor.detach().cpu().contiguous().view(torch.uint8).numpy().tobytes() + normalized = tensor.detach().cpu().contiguous() + if normalized.numel() == 0: + return b"" + # PyTorch cannot reinterpret a zero-dimensional tensor as a dtype with a + # different element size. Flattening preserves the exact storage bytes and + # also handles BF16, scalar adapter scales, and non-contiguous inputs. + return normalized.reshape(-1).view(torch.uint8).numpy().tobytes() def tensor_sha256(tensor: torch.Tensor) -> str: