hey for some reason i can no long run my multiple gpu train.py? It keeps filling up the first GPU rapidly, i thought maybe other gpus should also take some of the memory? Idk if the fsdp strategy is working? import os
import threading
import json
from fam.llm import training
import torch
import wandb
import numpy as np
from scipy.io import wavfile
from torch import nn
from torch.nn import functional as F
import pytorch_lightning as pl
from pytorch_lightning import Trainer
from pytorch_lightning.strategies import FSDPStrategy, DDPStrategy
from pytorch_lightning.plugins.environments import ClusterEnvironment
from huggingface_hub import snapshot_download, HfApi, HfFolder

# Specific to your project's structure
from fam.llm.training import parse_args, get_first_stage_path, get_second_stage_path, TrainingConfig, VoiceDataModule, WandbLogger, dist_utils, optimizer_utils, Evaluator, generate_audios, upload_file_async
from fam.llm.decoders import Decoder, EncodecDecoder
from fam.quantiser.text.tokenise import TrainedBPETokeniser
from fam.llm.adapters import FlattenedInterleavedEncodec2Codebook
from fam.llm.sample import Model
import datetime
import uuid

class MyClusterEnvironment(ClusterEnvironment):
    @property
    def creates_processes_externally(self) -> bool:
        """Return True if the cluster is managed (you don't launch processes yourself)"""
        return True

    def world_size(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_SIZE"])

    def global_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_RANK"])

    def local_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_LOCAL_RANK"])

    def node_rank(self) -> int:
        return int(os.environ["OMPI_COMM_WORLD_NODE_RANK"])

    main_address = os.getenv("MASTER_ADDR","")

    main_port = int(os.getenv("MASTER_PORT", "0"))
    
    def set_global_rank(self, rank):
        self.global_rank_ = rank

    def set_world_size(self, size):
        self.world_size_ = size

    def detect(self):
        return True

class KotobaSpeechModelFirstStage(pl.LightningModule):
    """
    A PyTorch Lightning module for the first stage of KotobaVoice model training.
    """
    
    def __init__(self, config_first_stage, config_second_stage, device, use_kv_cache, is_debug=False, configs=None):
        super().__init__()
        self.configs = vars(configs)
        self.config_first_stage = config_first_stage
        self.use_kv_cache = use_kv_cache
        self.prev_step = -1
        self.evaluator = Evaluator()
        self.train_evaluator = Evaluator()
        self.ignore_index = -1
        self.adj_block_size = config_first_stage.adj_block_size #Adjust for context length override else set to default block_size
        self.total_tokens = 0
        self.total_audio_tokens = 0
        self._wandb_initialized = False
        self._hf_initialized = False
        self._model_initialized = False

    def configure_model(self):
        """
        Configures the model and its components.
        """
        if self._model_initialized:
            return
        self.data_adapter = FlattenedInterleavedEncodec2Codebook(end_of_audio_token=1024)
        self.first_stage_model_cls = Model(
            self.config_first_stage,
            TrainedBPETokeniser,
            EncodecDecoder,
            data_adapter_fn=self.data_adapter.decode,
            use_kv_cache=self.use_kv_cache,
        )
        self.first_stage_model_transformer = self.first_stage_model_cls.model
        self._model_initialized = True

    def forward(self, text_tokens, embedding, inputs, targets):
        truncated_inputs = inputs[:,:self.adj_block_size] 
        truncated_first_targets = targets[:,:self.adj_block_size]
        truncated_inputs = truncated_inputs.unsqueeze(1)
        truncated_targets = truncated_first_targets.unsqueeze(1)
        # [batch_size, 7 (1 for first model, 6 for second), seq_len]

        list_logits, first_stage_unreduced_loss = self.first_stage_model_transformer(idx=truncated_inputs, embedding=embedding, targets=truncated_targets, loss_reduce="none")

        # Calculate the mean loss, excluding the ignored index (-1)
        # The operation inside the loss_reduce=mean calculation masks the -1 values 
        valid_mask = (truncated_targets != self.ignore_index)
        first_stage_loss = torch.mean(first_stage_unreduced_loss[valid_mask])

        # Calculate batch losses, excluding the ignored index for each batch
        first_stage_batch_losses = torch.tensor([
            torch.mean(first_stage_unreduced_loss[b][valid_mask[b]])
            for b in range(first_stage_unreduced_loss.size(0))
        ])

        # Asserting if the loss calculation above is the equivalent of PyTorch reduce.
        # _, assert_first_stage_loss = self.first_stage_model_transformer(idx=truncated_inputs, embedding=embedding, targets=truncated_targets, loss_reduce="mean")
        # assert(torch.isclose(first_stage_loss, assert_first_stage_loss, 1e-5))

        return first_stage_loss, first_stage_batch_losses

    def training_step(self, batch, batch_idx):
        tokens, embedding, audio_tokens, first_stage_input, first_stage_output, dataset_ids = batch 
        embedding = embedding.unsqueeze(1)

        loss, batch_losses = self(text_tokens=tokens, embedding=embedding, inputs=first_stage_input, targets=first_stage_output)

        if self.wandb_logger is not None:
            num_tokens = torch.sum(tokens != 0, 1)
            num_audio_tokens = torch.sum(first_stage_output != self.ignore_index, 1) + torch.sum(first_stage_output != self.ignore_index, 1)
            stats = {
                "dataset_id": dataset_ids,
                "num_tokens": num_tokens.tolist(),
                "num_audio_tokens": num_audio_tokens.tolist(),
                "loss": batch_losses.tolist(),
            }
            self.train_evaluator.update(stats)

        return loss
    
    def on_before_zero_grad(self, optimizer):
        if self.wandb_logger is not None:
            if self.global_step == 0:
                pass
            else:
                if dist_utils.is_main_process():
                    summary, dataset_summary = self.train_evaluator.summarize()
                    num_tokens = summary.pop("num_tokens")
                    num_audio_tokens = summary.pop("num_audio_tokens")
                    self.total_tokens += num_tokens
                    self.total_audio_tokens += num_audio_tokens

                    self.wandb_logger.log(results=summary, split="training", step=self.global_step, commit=False)
                    #self.wandb_logger.log(results=dataset_summary, split="training (details)", step=self.global_step, commit=False)
                    other_stats = {
                        f"lr_group{i}": list(self.optimizers().param_groups)[i]["lr"] for i in range(len(self.optimizers().param_groups))
                    }
                    other_stats["tokens"] = self.total_tokens
                    other_stats["audio_tokens"] = self.total_audio_tokens
                    other_stats["examples"] = self.configs["batch_size"] * self.global_step
                    self.wandb_logger.log(results=other_stats, split="stats", step=self.global_step, commit=False)

                self.train_evaluator.reset()


    def on_train_start(self): 
        if dist_utils.is_main_process():
            if self.configs["use_wandb"] and not self._wandb_initialized:
                self.wandb_logger = WandbLogger(training_args)
                wandb.watch(self.first_stage_model_transformer, log='parameters', log_freq=1000)
                self._wandb_initialized = True
            if self.configs["upload_checkpoint_to_hf"] and not self._hf_initialized:
                wandb_run_id = self.wandb_logger.get_run_id()
                self._hf_repo_name = "kotoba-speech/{}_{}".format(self.configs["exp_name"], wandb_run_id)
                api = HfApi()
                hf_token = os.environ.get('HF_TOKEN')
                api.create_repo(token=hf_token, repo_id=self._hf_repo_name, repo_type="model", exist_ok=True, private=True)
                self._hf_initialized = True
        else:
            self.wandb_logger = None
    
    def on_validation_start(self): 
        if dist_utils.is_main_process():
            if self.configs["use_wandb"] and not self._wandb_initialized:
                self.wandb_logger = WandbLogger(training_args)
                wandb.watch(self.first_stage_model_transformer, log='parameters', log_freq=1000)
                self._wandb_initialized = True
            if self.configs["upload_checkpoint_to_hf"] and not self._hf_initialized:
                wandb_run_id = self.wandb_logger.get_run_id()
                self._hf_repo_name = "kotoba-speech/{}_{}".format(self.configs["exp_name"], wandb_run_id)
                api = HfApi()
                hf_token = os.environ.get('HF_TOKEN')
                api.create_repo(token=hf_token, repo_id=self._hf_repo_name, repo_type="model", exist_ok=True, private=True)
                self._hf_initialized = True
        else:
            self.wandb_logger = None
            
    def validation_step(self, batch, batch_idx):
        tokens, embedding, audio_tokens, first_stage_input, first_stage_output, dataset_ids = batch
        tokens = tokens.unsqueeze(1)
        embedding = embedding.unsqueeze(1)
        loss, batch_losses = self(text_tokens=tokens, embedding=embedding, inputs=first_stage_input, targets=first_stage_output)

        stats = {
            "dataset_id": dataset_ids,
            "loss": batch_losses.tolist(),
        }
        self.evaluator.update(stats)
        return loss

    def on_validation_epoch_end(self):
        self.evaluator.synchronize_between_processes()
        if dist_utils.is_main_process():
            summary, dataset_summary = self.evaluator.summarize()
        self.evaluator.reset()
        if dist_utils.is_main_process():
            if hasattr(self, 'wandb_logger'):
                self.wandb_logger.log(results=summary, split="validation", step=self.global_step, commit=False)
                self.wandb_logger.log(results=dataset_summary, split="validation (details)", step=self.global_step, commit=False)    

    def on_validation_end(self):
        if dist_utils.is_main_process():
            if self.configs["upload_checkpoint_to_hf"]:
                if self.trainer.checkpoint_callback.last_model_path != "":
                    chpt_dir_name = self.trainer.checkpoint_callback.state_dict()["dirpath"]
                    files = [os.path.join(chpt_dir_name, f) for f in os.listdir(chpt_dir_name) if f.endswith(".ckpt") and f != "last.ckpt"]
                    newest_file = max(files, key=os.path.getmtime)
                    try:
                        upload_thread = threading.Thread(target=upload_file_async, args=(newest_file, os.path.basename(newest_file), self._hf_repo_name, "model"))
                        upload_thread.start()
                    except:
                        print("HF upload failed for checkpoint: {}".format(newest_file))

    def on_save_checkpoint(self, checkpoint):
        if dist_utils.is_main_process():
            # TODO: Fix this. This is probably saving checkpoints twice. Ideally, overwriting on_save_checkpoint in the callback is cleaner.
            if self.configs["log_audio_sample"]:
                with torch.no_grad():
                    ckpt_path = os.path.join("lightning_logs", uuid.uuid4().hex + ".pt")
                    torch.save(checkpoint, ckpt_path)
                    print("Generating sample audios for checkpoint : {}".format(ckpt_path))
                    results = generate_audios(first_model_path=ckpt_path, file_path=self.configs["validation_file_path"], guidance_scale=self.configs["guidance_scale"])
                    audio_files = {}
                    for i, rs in enumerate(results):
                        id = rs["id"] if "id" in rs else "ex-{}".format(i)
                        caption = rs["prompt"]
                        wavpath = rs["audio_path"]
                        sample_rate, data = wavfile.read(wavpath)
                        if data.dtype != np.float32:
                            data = data.astype(np.float32) / np.max(np.abs(data))  # Normalize data to [-1, 1]
                        audio_files[id] = wandb.Audio(data, caption=caption, sample_rate=sample_rate)
                    self.wandb_logger.log(results=audio_files, split="audio", step=self.global_step, commit=False)
                    os.remove(ckpt_path)

    def configure_optimizers(self):
        return optimizer_utils.set_schedule(self)

# Train the model
if __name__ == "__main__":
    training_args = parse_args()
    data_module = VoiceDataModule(training_args.per_gpu_batchsize, training_args.data_list, training_args.local_path, training_args.ratio_filter, training_args.per_video_spk)

    # assertions
    if training_args.upload_checkpoint_to_hf:
        assert training_args.use_wandb, "HF uploading requires wandb run_id, so use_wandb must be enabled."
        # Validate HF token
        api = HfApi()
        try:
            hf_token = os.environ.get('HF_TOKEN')
            if not hf_token:
                raise ValueError("HF_TOKEN environment variable is not set or empty. Please export HF_TOKEN by running 'export HF_TOKEN=YOUR_TOKEN'.")

            user_info = api.whoami(token=hf_token)
            print("User info retrieved:", user_info)
        except Exception as e:
            raise ValueError("Invalid HF token provided or error retrieving user info. Please ensure the HF_TOKEN is correctly set.") from e


    model_dir = snapshot_download(repo_id=training_args.huggingface_repo_id)
    first_stage_ckpt_path = get_first_stage_path(model_dir)
    second_stage_ckpt_path = get_second_stage_path(model_dir)
    config_first_stage = TrainingConfig(
        ckpt_path=first_stage_ckpt_path,
        num_samples=training_args.num_samples,
        seed=training_args.seed,
        device=training_args.device,
        dtype=training_args.dtype,
        compile=training_args.compile,
        init_from=training_args.init_from,
        train_from_scratch=training_args.train_from_scratch,
        output_dir=training_args.output_dir,
        spkemb_dropout=training_args.spkemb_dropout,
        block_size=training_args.block_size,
        adj_block_size=training_args.block_size if training_args.block_size_override is None else training_args.block_size_override, #linear interpolates existing positional embeddings to extend CL.
        adj_ckpt_path=training_args.ckpt_path #This will override the default ckpt path if specified as an arg.
    )

    config_second_stage = TrainingConfig(
        ckpt_path=second_stage_ckpt_path,
        num_samples=training_args.num_samples,
        seed=training_args.seed,
        device=training_args.device,
        dtype=training_args.dtype,
        compile=training_args.compile,
        init_from=training_args.init_from,
        train_from_scratch=training_args.train_from_scratch,
        output_dir=training_args.output_dir,
        spkemb_dropout=training_args.spkemb_dropout,
        block_size=training_args.block_size,
        adj_block_size=training_args.block_size if training_args.block_size_override is None else training_args.block_size_override,
        adj_ckpt_path=training_args.ckpt_path #This will override the default ckpt path if specified as an arg.
    )

    is_debug = training_args.debug
    model = KotobaSpeechModelFirstStage(config_first_stage, config_second_stage, device=training_args.device, use_kv_cache=training_args.use_kv_cache, is_debug=is_debug, configs=training_args)
    model.configure_model()

    # # #

    #finetune loading

    # # #

    if training_args.finetuning_ckpt_path is not None:
        ckpt = torch.load(training_args.finetuning_ckpt_path)
        state_dict = ckpt["state_dict"]
        model.load_state_dict(state_dict)
    lr_callback = pl.callbacks.LearningRateMonitor(logging_interval="step")

  
  
  
    if training_args.savedir_path is not None:
        save_path = training_args.savedir_path
    else:
        current_time = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        save_path = "/data/checkpoints/{}".format(current_time)

    checkpoint_callback = pl.callbacks.ModelCheckpoint(
        dirpath=save_path,
        save_top_k=1,
        save_last=True,
        verbose=True,
        monitor=None,
        #save_on_train_epoch_end=True,
    )

    if training_args.log_audio_sample:
        assert checkpoint_callback.save_last,  "last_model_path will be empty and code will break."
    callbacks = [checkpoint_callback]

    num_gpus = (
        training_args.num_gpus
        if isinstance(training_args.num_gpus, int)
        else len(training_args.num_gpus)
    )

    grad_steps = training_args.batch_size // (
        training_args.per_gpu_batchsize * num_gpus * training_args.num_nodes
    )

    max_steps = training_args.max_steps if training_args.max_steps is not None else None
    
    assert training_args.dtype =='bfloat16'
    if training_args.dtype == "bfloat16":
        precision="bf16-mixed"
    else:
        raise "Precision needs to be studied well."
    
    assert training_args.fsdp_strategy is not None
    if training_args.fsdp_strategy is not None: 
        strategy = FSDPStrategy(
            sharding_strategy=training_args.fsdp_strategy
        )
    elif training_args.use_ddp_strategy: 
        strategy = DDPStrategy(
            static_graph=True,
        )
    else:
        strategy = 'ddp_find_unused_parameters_true'
    
    if training_args.num_nodes > 1:
        plugins = [MyClusterEnvironment()]
    else:
        plugins = None
    
    if training_args.val_check_interval is not None:
        if training_args.val_check_interval >= 1:
            training_args.val_check_interval=int(training_args.val_check_interval)


    trainer = Trainer(
        callbacks=callbacks,
        devices=training_args.num_gpus,
        strategy=strategy,
        num_nodes=training_args.num_nodes,
        precision=precision,
        accelerator="cuda",
        benchmark=True,
        deterministic=True,
        max_epochs=training_args.max_epoch if max_steps is None else 1000,
        accumulate_grad_batches=grad_steps,
        val_check_interval=training_args.val_check_interval,
        check_val_every_n_epoch=training_args.check_val_every_n_epoch,
        log_every_n_steps=10,
        fast_dev_run=training_args.fast_dev_run,
        plugins=plugins,
        gradient_clip_val=training_args.gradient_clip_val, 
    )

    if training_args.finetuning_ckpt_path is not None:
        training_args.ckpt_path = None

    #If we have block override do not load the original ckpt through lightning (its manually loaded above in the Model class)
    #The lightning module ckpt loading does not support the needed embedding interpolations.
    ckpt_path = training_args.ckpt_path if training_args.block_size_override is None else None

    trainer.fit(
        model,        
        ckpt_path=ckpt_path,
        datamodule=data_module
    )
    trainer.print(torch.cuda.memory_summary())