Training LoRA using FSDP config

[hannalaguilar]

Hi, thanks a lot for this project!

I’m trying to train a LoRA model using your FSDP config (configs/accelerate/fsdp.yaml) with num_processes=4 on a machine with 2 H100 GPUs (80 GB VRAM each).

I launch training with:

CUDA_VISIBLE_DEVICES=0,1 \
uv run accelerate launch --config_file configs/accelerate/fsdp.yaml \
  scripts/train.py configs/ltx2_ti2v_lora.yaml

During training startup, I get the following error on rank1:

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 30.00 MiB. GPU 0 has a total capacity of 79.10 GiB of which 15.19 MiB is free. Process 16746 has 20.84 GiB memory in use. Process 16747 has 20.33 GiB memory in use. Process 16748 has 18.94 GiB memory in use. Process 16749 has 18.94 GiB memory in use. Of the allocated memory 19.60 GiB is allocated by PyTorch, and 225.59 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

The error is here:

self._cached_validation_embeddings = self._load_text_encoder_and_cache_embeddings()

When using num_processes=4, it seems that the text encoder is loaded on GPU 0 by all processes.

My config:

[jiaxiangc]

same issue + 1.
We find that the gpu memory in rank 0 is easily OOM. And other ranks only 0 GB.

[vasileone]

same issue + 1.
Hope to fix this issue as soon as possible.

[rex-29]

It seems like the whole project isn’t ready yet.

[rex-29]

Change the "_load_text_encoder_and_cache_embeddings" in LTX-2/packages/ltx-trainer/src/ltx_trainer/trainer.py

def _load_text_encoder_and_cache_embeddings(self) -> list[CachedPromptEmbeddings] | None:
        """Load text encoder, computes and returns validation embeddings."""

        # This method:
        #   1. Loads the text encoder on GPU
        #   2. If validation prompts are configured, computes and caches their embeddings
        #   3. Unloads the heavy Gemma model while keeping the lightweight embedding connectors
        #   The text encoder is kept (as self._text_encoder) but with model/tokenizer/feature_extractor
        #   set to None. Only the embedding connectors remain for use during training.
        local_rank = int(os.environ.get("LOCAL_RANK", 0))
        device = f"cuda:{local_rank}"
        # Load text encoder on GPU
        logger.debug("Loading text encoder...")
        if self._config.acceleration.load_text_encoder_in_8bit:
            logger.warning(
                "⚠️  load_text_encoder_in_8bit is set to True but 8-bit text encoder loading "
                "is not currently implemented. The text encoder will be loaded in bfloat16 precision."
            )
        self._text_encoder = load_text_encoder(
            checkpoint_path=self._config.model.model_path,
            gemma_model_path=self._config.model.text_encoder_path,
            device=device,
            dtype=torch.bfloat16,
        )

        # Cache validation embeddings if prompts are configured
        cached_embeddings = None
        if self._config.validation.prompts:
            logger.info(f"Pre-computing embeddings for {len(self._config.validation.prompts)} validation prompts...")
            cached_embeddings = []
            with torch.inference_mode():
                for prompt in self._config.validation.prompts:
                    v_ctx_pos, a_ctx_pos, _ = self._text_encoder(prompt)
                    v_ctx_neg, a_ctx_neg, _ = self._text_encoder(self._config.validation.negative_prompt)

                    cached_embeddings.append(
                        CachedPromptEmbeddings(
                            video_context_positive=v_ctx_pos.cpu(),
                            audio_context_positive=a_ctx_pos.cpu(),
                            video_context_negative=v_ctx_neg.cpu() if v_ctx_neg is not None else None,
                            audio_context_negative=a_ctx_neg.cpu() if a_ctx_neg is not None else None,
                        )
                    )

        # Unload heavy components to free VRAM, keeping only the embedding connectors
        self._text_encoder.model = None
        self._text_encoder.tokenizer = None
        self._text_encoder.feature_extractor_linear = None
        torch.cuda.empty_cache()

        logger.debug("Validation prompt embeddings cached. Gemma model unloaded")
        return cached_embeddings

works for me.

[LittleNyima]

Fixed this in #67