Add files using upload-large-folder tool

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_checkpointing.py

@@ -0,0 +1,269 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import json
+import os
+
+import evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+from accelerate import Accelerator, DistributedType
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.XLA:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def evaluation_loop(accelerator, model, eval_dataloader, metric):
+    model.eval()
+    samples_seen = 0
+    for step, batch in enumerate(eval_dataloader):
+        # We could avoid this line since we set the accelerator with `device_placement=True`.
+        batch.to(accelerator.device)
+        with torch.no_grad():
+            outputs = model(**batch)
+        predictions = outputs.logits.argmax(dim=-1)
+        # It is slightly faster to call this once, than multiple times
+        predictions, references = accelerator.gather(
+            (predictions, batch["labels"])
+        )  # If we are in a multiprocess environment, the last batch has duplicates
+        if accelerator.use_distributed:
+            if step == len(eval_dataloader) - 1:
+                predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                references = references[: len(eval_dataloader.dataset) - samples_seen]
+            else:
+                samples_seen += references.shape[0]
+        metric.add_batch(
+            predictions=predictions,
+            references=references,
+        )
+
+    eval_metric = metric.compute()
+    return eval_metric["accuracy"]
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+    else:
+        gradient_accumulation_steps = 1
+    max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps
+
+    # Instantiate scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+    metric = evaluate.load("glue", "mrpc")
+    ending_epoch = num_epochs
+
+    if args.partial_train_epoch is not None:
+        ending_epoch = args.partial_train_epoch
+
+    if args.resume_from_checkpoint:
+        accelerator.load_state(args.resume_from_checkpoint)
+        epoch_string = args.resume_from_checkpoint.split("epoch_")[1]
+        state_epoch_num = ""
+        for char in epoch_string:
+            if char.isdigit():
+                state_epoch_num += char
+            else:
+                break
+        starting_epoch = int(state_epoch_num) + 1
+        accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric)
+        accelerator.print("resumed checkpoint performance:", accuracy)
+        accelerator.print("resumed checkpoint's scheduler's lr:", lr_scheduler.get_lr()[0])
+        accelerator.print("resumed optimizers's lr:", optimizer.param_groups[0]["lr"])
+        with open(os.path.join(args.output_dir, f"state_{starting_epoch - 1}.json")) as f:
+            resumed_state = json.load(f)
+            assert resumed_state["accuracy"] == accuracy, "Accuracy mismatch, loading from checkpoint failed"
+            assert resumed_state["lr"] == lr_scheduler.get_lr()[0], (
+                "Scheduler learning rate mismatch, loading from checkpoint failed"
+            )
+            assert resumed_state["optimizer_lr"] == optimizer.param_groups[0]["lr"], (
+                "Optimizer learning rate mismatch, loading from checkpoint failed"
+            )
+            assert resumed_state["epoch"] == starting_epoch - 1, "Epoch mismatch, loading from checkpoint failed"
+            return
+
+    # Now we train the model
+    state = {}
+    for epoch in range(starting_epoch, ending_epoch):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            outputs = model(**batch)
+            loss = outputs.loss
+            loss = loss / gradient_accumulation_steps
+            accelerator.backward(loss)
+            if step % gradient_accumulation_steps == 0:
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+            overall_step += 1
+        output_dir = f"epoch_{epoch}"
+        output_dir = os.path.join(args.output_dir, output_dir)
+        accelerator.save_state(output_dir)
+        accuracy = evaluation_loop(accelerator, model, eval_dataloader, metric)
+        state["accuracy"] = accuracy
+        state["lr"] = lr_scheduler.get_lr()[0]
+        state["optimizer_lr"] = optimizer.param_groups[0]["lr"]
+        state["epoch"] = epoch
+        state["step"] = overall_step
+        accelerator.print(f"epoch {epoch}:", state)
+
+        accelerator.wait_for_everyone()
+        if accelerator.is_main_process:
+            with open(os.path.join(args.output_dir, f"state_{epoch}.json"), "w") as f:
+                json.dump(state, f)
+    accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--resume_from_checkpoint",
+        type=str,
+        default=None,
+        help="If the training should continue from a checkpoint folder.",
+    )
+    parser.add_argument(
+        "--partial_train_epoch",
+        type=int,
+        default=None,
+        help="If passed, the training will stop after this number of epochs.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=2,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_alst_ulysses_sp.py

@@ -0,0 +1,129 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Test script for verifying ALST/Ulysses SP works
+"""
+
+import torch
+from deepspeed.runtime.utils import move_to_device
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from accelerate import Accelerator
+from accelerate.utils import ParallelismConfig, set_seed
+from accelerate.utils.dataclasses import DeepSpeedSequenceParallelConfig
+
+
+set_seed(42)
+
+world_size = 2
+model_name = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+
+micro_batch_size = 1
+
+parallelism_config = ParallelismConfig(
+    sp_backend="deepspeed",
+    sp_size=world_size,
+    # dp_shard_size=1, # set if dp is wanted as well
+    sp_handler=DeepSpeedSequenceParallelConfig(
+        sp_seq_length=256,
+        sp_seq_length_is_variable=True,
+        sp_attn_implementation="sdpa",
+    ),
+)
+
+accelerator = Accelerator(
+    parallelism_config=parallelism_config,
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
+samples = 4
+seqlen = 32
+input_ids = torch.arange(1, seqlen * samples + 1).view(-1, seqlen) + 100
+position_ids = torch.arange(seqlen * samples).view(-1, seqlen)
+
+ds = torch.utils.data.TensorDataset(input_ids, position_ids)
+
+
+def collate_fn(batch):
+    input_ids, position_ids = batch[0]
+    return dict(
+        input_ids=input_ids.unsqueeze(0),
+        position_ids=position_ids.unsqueeze(0),
+        labels=input_ids.unsqueeze(0),
+    )
+
+
+dl = torch.utils.data.DataLoader(ds, batch_size=micro_batch_size, collate_fn=collate_fn)
+
+optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
+
+rank = torch.distributed.get_rank()
+
+if rank == 0:
+    print(f"DL orig: {len(dl)} samples")
+
+model, optimizer, dl = accelerator.prepare(model, optimizer, dl)
+
+if rank == 0:
+    print(f"DL w/ adapter: {len(dl)} samples")
+
+sp_size = parallelism_config.sp_size if parallelism_config else 1
+if sp_size > 1:
+    sp_group = accelerator.torch_device_mesh["sp"].get_group()
+    sp_world_size = parallelism_config.sp_size
+
+unwrapped_model = accelerator.unwrap_model(model)
+
+# Normal training loop
+for iter, batch in enumerate(dl):
+    optimizer.zero_grad()
+
+    if rank == 0:
+        print(f"batch {iter}: seqlen: {len(batch['input_ids'][0])}")
+    batch = move_to_device(batch, model.device)
+    outputs = model(**batch)
+
+    shift_labels = batch["shift_labels"]
+    loss = unwrapped_model.loss_function(
+        logits=outputs.logits,
+        labels=None,
+        shift_labels=shift_labels,
+        vocab_size=unwrapped_model.config.vocab_size,
+    )
+
+    if sp_size > 1:
+        # differentiable weighted per-shard-loss aggregation across ranks
+        losses_per_rank = torch.distributed.nn.functional.all_gather(loss, group=sp_group)
+        # special dealing with SFT that has prompt tokens that aren't used in loss computation
+        good_tokens = (shift_labels != -100).view(-1).sum()
+        good_tokens_per_rank = torch.distributed.nn.functional.all_gather(good_tokens, group=sp_group)
+        total_loss = sum(
+            losses_per_rank[rank] * good_tokens_per_rank[rank]
+            for rank in range(sp_world_size)
+            if good_tokens_per_rank[rank] > 0
+        )
+        total_good_tokens = sum(good_tokens_per_rank)
+        loss = total_loss / max(total_good_tokens, 1)
+
+    if rank == 0:
+        accelerator.print(f"{iter}: {loss=}")
+    accelerator.log(dict(train_loss=loss, step=iter))
+
+    accelerator.backward(loss)
+    optimizer.step()
+
+accelerator.end_training()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_ds_multiple_model.py

@@ -0,0 +1,331 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Test script for verifying multiple models can be utilized with Accelerate + DeepSpeed:
+
+Scenario 1: One model is training, another model is being used for inference/logits to impact training in some form.
+Scenario 2: Two models are training simultaneously, which means two optimizers, etc.
+"""
+
+import argparse
+from pathlib import Path
+
+import evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup
+
+from accelerate import Accelerator, DeepSpeedPlugin, DistributedType
+from accelerate.state import AcceleratorState
+from accelerate.utils.deepspeed import get_active_deepspeed_plugin
+
+
+EVAL_BATCH_SIZE = 16
+
+
+class NoiseModel(torch.nn.Module):
+    def __init__(self, noise_factor=0.1):
+        super().__init__()
+        self.noise_factor = torch.nn.Parameter(torch.tensor(noise_factor, dtype=torch.float32))
+
+    def forward(self, loss):
+        return loss * self.noise_factor
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.XLA:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+test_file_path = __file__
+path = Path(test_file_path).resolve()
+test_file_dir_str = str(path.parent.parent.parent.parent.parent.parent)
+
+# Create our DS plugins
+# We use custom schedulers and optimizers, hence `model_only`
+ds_config_file = dict(
+    zero2=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero2_model_only.json",
+    zero3=f"{test_file_dir_str}/tests/deepspeed/ds_config_zero3_model_only.json",
+)
+
+
+def single_model_training(config, args):
+    # Training a single model, we have a `noise` model that is untrainable used to inject some noise into the training process
+    num_epochs = config["num_epochs"]
+    zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"])
+    zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"])
+
+    deepspeed_plugins = {"training": zero2_plugin, "inference": zero3_plugin}
+
+    # Initialize accelerator
+    accelerator = Accelerator(
+        deepspeed_plugins=deepspeed_plugins,
+        mixed_precision="bf16",
+    )
+
+    # Initialize model under zero2 plugin
+    assert get_active_deepspeed_plugin(accelerator.state) is zero2_plugin
+    train_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
+    train_dataloader, eval_dataloader = get_dataloaders(
+        accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path
+    )
+    max_training_steps = len(train_dataloader) * config["num_epochs"]
+    optimizer = AdamW(train_model.parameters(), lr=config["lr"])
+    lr_scheduler = get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=0, num_training_steps=max_training_steps
+    )
+
+    train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler = accelerator.prepare(
+        train_dataloader, eval_dataloader, train_model, optimizer, lr_scheduler
+    )
+
+    # Now prepare the model under zero3 plugin
+    accelerator.state.select_deepspeed_plugin("inference")
+    assert get_active_deepspeed_plugin(accelerator.state) is zero3_plugin
+    inference_model = NoiseModel()
+    inference_model = accelerator.prepare(inference_model)
+    inference_model.eval()
+
+    # Run training loop
+    accelerator.state.select_deepspeed_plugin("training")
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Now we train the model
+    best_performance = 0
+    metric = evaluate.load("glue", "mrpc")
+    performance_metric = {}
+    for epoch in range(starting_epoch, num_epochs):
+        train_model.train()
+        inference_model.train()
+        for step, batch in enumerate(train_dataloader):
+            with accelerator.accumulate(train_model):
+                outputs_1 = train_model(**batch)
+                with torch.no_grad():
+                    outputs_2 = inference_model(outputs_1.loss)
+                # Combine the losses
+                loss = outputs_1.loss + outputs_2
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+        train_model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            with torch.no_grad():
+                outputs = train_model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather_for_metrics((predictions, batch["labels"]))
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+        performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"]
+
+        if best_performance < eval_metric["accuracy"]:
+            best_performance = eval_metric["accuracy"]
+    assert best_performance > performance_metric["epoch-0"]
+
+
+def multiple_model_training(config, args):
+    # This will essentially be like a k-fold model, but one model is Zero-2 and another model is Zero-3
+    num_epochs = config["num_epochs"]
+    zero2_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero2"])
+    zero3_plugin = DeepSpeedPlugin(hf_ds_config=ds_config_file["zero3"])
+
+    deepspeed_plugins = {"zero2": zero2_plugin, "zero3": zero3_plugin}
+
+    # Initialize accelerator
+    zero2_accelerator = Accelerator(
+        deepspeed_plugins=deepspeed_plugins,
+        mixed_precision="bf16",
+    )
+
+    # Since an `AcceleratorState` has already been made, we can just reuse it here
+    zero3_accelerator = Accelerator()
+
+    # Initialize model under zero2 plugin
+    assert get_active_deepspeed_plugin(zero2_accelerator.state) is zero2_plugin
+    zero2_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
+    train_dataloader, eval_dataloader = get_dataloaders(
+        zero2_accelerator, batch_size=config["batch_size"], model_name=args.model_name_or_path
+    )
+    max_training_steps = len(train_dataloader) * config["num_epochs"]
+    zero2_optimizer = AdamW(zero2_model.parameters(), lr=config["lr"])
+    zero2_lr_scheduler = get_linear_schedule_with_warmup(
+        zero2_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps
+    )
+
+    train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler = zero2_accelerator.prepare(
+        train_dataloader, eval_dataloader, zero2_model, zero2_optimizer, zero2_lr_scheduler
+    )
+    assert zero2_accelerator.deepspeed_engine_wrapped.engine is zero2_model
+
+    # now do Zero3
+    zero3_accelerator.state.select_deepspeed_plugin("zero3")
+    zero3_plugin.deepspeed_config["train_micro_batch_size_per_gpu"] = zero2_plugin.deepspeed_config[
+        "train_micro_batch_size_per_gpu"
+    ]
+    assert get_active_deepspeed_plugin(zero3_accelerator.state) is zero3_plugin
+    zero3_model = AutoModelForSequenceClassification.from_pretrained(args.model_name_or_path)
+    zero3_optimizer = AdamW(zero3_model.parameters(), lr=config["lr"])
+    zero3_lr_scheduler = get_linear_schedule_with_warmup(
+        zero3_optimizer, num_warmup_steps=0, num_training_steps=max_training_steps
+    )
+    zero3_model, zero3_optimizer, zero3_lr_scheduler = zero3_accelerator.prepare(
+        zero3_model, zero3_optimizer, zero3_lr_scheduler
+    )
+    assert zero3_accelerator.deepspeed_engine_wrapped.engine is zero3_model
+
+    # Run training loop
+    starting_epoch = 0
+
+    # Now we train the model
+    best_performance_a = 0
+    best_performance_b = 0
+    metric_a = evaluate.load("glue", "mrpc")
+    metric_b = evaluate.load("glue", "mrpc")
+    performance_metric_a = {}
+    performance_metric_b = {}
+    for epoch in range(starting_epoch, num_epochs):
+        zero2_model.train()
+        zero3_model.train()
+        for step, batch in enumerate(train_dataloader):
+            with zero2_accelerator.accumulate(zero2_model, zero3_model):
+                outputs_1 = zero2_model(**batch)
+                zero2_accelerator.backward(outputs_1.loss)
+                zero2_optimizer.step()
+                zero2_lr_scheduler.step()
+                zero2_optimizer.zero_grad()
+                outputs_2 = zero3_model(**batch)
+                zero3_accelerator.backward(outputs_2.loss)
+                zero3_optimizer.step()
+                zero3_lr_scheduler.step()
+                zero3_optimizer.zero_grad()
+
+        zero2_model.eval()
+        zero3_model.eval()
+        for step, batch in enumerate(eval_dataloader):
+            with torch.no_grad():
+                logits_a = zero2_model(**batch).logits
+                logits_b = zero3_model(**batch).logits
+            # Combine the logits from both models
+            predictions_a = logits_a.argmax(dim=-1)
+            predictions_b = logits_b.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions_a, predictions_b, references = zero2_accelerator.gather_for_metrics(
+                (predictions_a, predictions_b, batch["labels"])
+            )
+            metric_a.add_batch(
+                predictions=predictions_a,
+                references=references,
+            )
+            metric_b.add_batch(
+                predictions=predictions_b,
+                references=references,
+            )
+
+        eval_metric_a = metric_a.compute()
+        eval_metric_b = metric_b.compute()
+        # Use accelerator.print to print only on the main process.
+        zero2_accelerator.print(f"epoch {epoch}:", eval_metric_a, eval_metric_b)
+        performance_metric_a[f"epoch-{epoch}"] = eval_metric_a["accuracy"]
+        performance_metric_b[f"epoch-{epoch}"] = eval_metric_b["accuracy"]
+
+        if best_performance_a < eval_metric_a["accuracy"]:
+            best_performance_a = eval_metric_a["accuracy"]
+        if best_performance_b < eval_metric_b["accuracy"]:
+            best_performance_b = eval_metric_b["accuracy"]
+    assert best_performance_a > performance_metric_a["epoch-0"]
+    assert best_performance_b > performance_metric_b["epoch-0"]
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--performance_lower_bound",
+        type=float,
+        default=None,
+        help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=3,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 8}
+    single_model_training(config, args)
+    AcceleratorState._reset_state(True)
+    multiple_model_training(config, args)
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_metrics.py

@@ -0,0 +1,307 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+import math
+import os
+from copy import deepcopy
+
+import datasets
+import evaluate
+import torch
+import transformers
+from datasets import load_dataset
+from torch.utils.data import DataLoader, IterableDataset
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+
+from accelerate import Accelerator, DataLoaderConfiguration, DistributedType
+from accelerate.data_loader import DataLoaderDispatcher
+from accelerate.test_utils import RegressionDataset, RegressionModel, torch_device
+from accelerate.utils import is_torch_xla_available, set_seed
+
+
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "true"
+
+
+class ListHandler(logging.Handler):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logs = []
+
+    def emit(self, record):
+        self.logs.append(record)
+
+
+def get_basic_setup(accelerator, num_samples=82, batch_size=16):
+    "Returns everything needed to perform basic training"
+    set_seed(42)
+    model = RegressionModel()
+    ddp_model = deepcopy(model)
+    dset = RegressionDataset(length=num_samples)
+    dataloader = DataLoader(dset, batch_size=batch_size)
+    model.to(accelerator.device)
+    ddp_model, dataloader = accelerator.prepare(ddp_model, dataloader)
+    return model, ddp_model, dataloader
+
+
+def get_dataloader(accelerator: Accelerator, use_longest=False):
+    tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/mrpc-bert-base-cased")
+    dataset = load_dataset("glue", "mrpc", split="validation")
+
+    def tokenize_function(examples):
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    with accelerator.main_process_first():
+        tokenized_datasets = dataset.map(
+            tokenize_function,
+            batched=True,
+            remove_columns=["idx", "sentence1", "sentence2"],
+        )
+
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        if use_longest:
+            return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+        return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+
+    return DataLoader(tokenized_datasets, shuffle=False, collate_fn=collate_fn, batch_size=16)
+
+
+def get_mrpc_setup(dispatch_batches, split_batches):
+    dataloader_config = DataLoaderConfiguration(dispatch_batches=dispatch_batches, split_batches=split_batches)
+    accelerator = Accelerator(dataloader_config=dataloader_config)
+    dataloader = get_dataloader(accelerator, not dispatch_batches)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        "hf-internal-testing/mrpc-bert-base-cased", return_dict=True
+    )
+    ddp_model, ddp_dataloader = accelerator.prepare(model, dataloader)
+    return {
+        "ddp": [ddp_model, ddp_dataloader, torch_device],
+        "no": [model, dataloader, accelerator.device],
+    }, accelerator
+
+
+def generate_predictions(model, dataloader, accelerator):
+    logits_and_targets = []
+    for batch in dataloader:
+        input, target = batch.values()
+        with torch.no_grad():
+            logit = model(input)
+            logit, target = accelerator.gather_for_metrics((logit, target))
+            logits_and_targets.append((logit, target))
+    logits, targs = [], []
+    for logit, targ in logits_and_targets:
+        logits.append(logit)
+        targs.append(targ)
+    logits, targs = torch.cat(logits), torch.cat(targs)
+    return logits, targs
+
+
+def test_torch_metrics(
+    accelerator: Accelerator, num_samples=82, dispatch_batches=False, split_batches=False, batch_size=16
+):
+    _, ddp_model, dataloader = get_basic_setup(accelerator, num_samples, batch_size)
+    logits, _ = generate_predictions(ddp_model, dataloader, accelerator)
+    assert len(logits) == num_samples, (
+        f"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(logits)}"
+    )
+
+
+def test_mrpc(dispatch_batches: bool = False, split_batches: bool = False):
+    metric = evaluate.load("glue", "mrpc")
+    setup, accelerator = get_mrpc_setup(dispatch_batches, split_batches)
+    # First do baseline
+    model, dataloader, device = setup["no"]
+    model.to(device)
+    model.eval()
+    for batch in dataloader:
+        batch.to(device)
+        with torch.inference_mode():
+            outputs = model(**batch)
+        preds = outputs.logits.argmax(dim=-1)
+        metric.add_batch(predictions=preds, references=batch["labels"])
+    baseline = metric.compute()
+
+    # Then do distributed
+    model, dataloader, device = setup["ddp"]
+    model.eval()
+    for batch in dataloader:
+        with torch.inference_mode():
+            outputs = model(**batch)
+        preds = outputs.logits.argmax(dim=-1)
+        references = batch["labels"]
+        preds, references = accelerator.gather_for_metrics((preds, references))
+        metric.add_batch(predictions=preds, references=references)
+    distributed = metric.compute()
+
+    for key in "accuracy f1".split():
+        assert math.isclose(baseline[key], distributed[key]), (
+            f"Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n"
+        )
+
+
+def test_gather_for_metrics_with_non_tensor_objects_iterable_dataset():
+    class DummyIterableDataset(IterableDataset):
+        def __init__(self, data):
+            self.data = data
+
+        def __len__(self):
+            return len(self.data)
+
+        def __iter__(self):
+            yield from self.data
+
+    iterable_dataset = DummyIterableDataset([n for n in range(30)])
+    dataloader = DataLoader(iterable_dataset, batch_size=4)
+    accelerator = Accelerator()
+    prepared_dataloader = accelerator.prepare(dataloader)
+
+    if accelerator.is_main_process:
+        logger = logging.root.manager.loggerDict["accelerate.accelerator"]
+        list_handler = ListHandler()
+        logger.addHandler(list_handler)
+
+    batches_for_metrics = []
+    for batch in prepared_dataloader:
+        batches_for_metrics.append(accelerator.gather_for_metrics(batch))
+
+    assert torch.cat(batches_for_metrics).size(0) == 30
+
+    if accelerator.is_main_process:
+        assert len(list_handler.logs) == 0
+        logger.removeHandler(list_handler)
+
+
+def test_gather_for_metrics_with_iterable_dataset():
+    class DummyIterableDataset(IterableDataset):
+        def __init__(self, data):
+            self.data = data
+
+        def __len__(self):
+            return len(self.data)
+
+        def __iter__(self):
+            yield from self.data
+
+    iterable_dataset = DummyIterableDataset(torch.as_tensor(range(30)))
+    dataloader = DataLoader(iterable_dataset, batch_size=4)
+
+    accelerator = Accelerator()
+    prepared_dataloader = accelerator.prepare(dataloader)
+
+    assert isinstance(prepared_dataloader, DataLoaderDispatcher)
+
+    if accelerator.is_main_process:
+        logger = logging.root.manager.loggerDict["accelerate.accelerator"]
+        list_handler = ListHandler()
+        logger.addHandler(list_handler)
+
+    batches_for_metrics = []
+    for batch in prepared_dataloader:
+        batches_for_metrics.append(accelerator.gather_for_metrics(batch))
+
+    assert torch.cat(batches_for_metrics).size(0) == 30
+
+    if accelerator.is_main_process:
+        assert len(list_handler.logs) == 0
+
+        logger.removeHandler(list_handler)
+
+
+def test_gather_for_metrics_drop_last():
+    accelerator = Accelerator()
+    per_device_batch_size = 5
+    num_items = (10 * accelerator.num_processes) + 1
+    dataloader = DataLoader(range(num_items), batch_size=per_device_batch_size, drop_last=True)
+    dataloader = accelerator.prepare(dataloader)
+
+    iterator = iter(dataloader)
+    next(iterator)  # Skip first batch tensor([0, 1, 2, 3, 4], device='cuda:0')
+    batch = next(iterator)
+    gathered_items = accelerator.gather_for_metrics(batch)
+
+    # Should return a full set of complete batches from each GPU
+    num_expected_items = per_device_batch_size * accelerator.num_processes
+    assert gathered_items.size(0) == (num_expected_items), (
+        f"Expected number of items: {num_expected_items}, Actual: {gathered_items.size(0)}"
+    )
+
+
+def main():
+    dataloader_config = DataLoaderConfiguration(split_batches=False, dispatch_batches=False)
+    accelerator = Accelerator(dataloader_config=dataloader_config)
+    if accelerator.is_local_main_process:
+        datasets.utils.logging.set_verbosity_warning()
+        transformers.utils.logging.set_verbosity_warning()
+    else:
+        datasets.utils.logging.set_verbosity_error()
+        transformers.utils.logging.set_verbosity_error()
+    # TorchXLA does not support batch dispatching. 'put_on_device' is always False for
+    # TorchXLA, which can cause a value error in 'prepare_data_loader' function.
+    dispatch_batches_options = [False] if accelerator.state.distributed_type == DistributedType.XLA else [True, False]
+
+    # Temporarily close this test for TorchXLA due to the 'Cannot set version_counter for
+    # inference tensor' error in inference mode. Reopen it after TorchXLA fixes this bug.
+    # These are a bit slower so they should only be ran on the GPU or TPU
+    if accelerator.device.type != "cpu" and not is_torch_xla_available():
+        if accelerator.is_local_main_process:
+            print("**Testing gather_for_metrics**")
+        for split_batches in [True, False]:
+            for dispatch_batches in dispatch_batches_options:
+                if accelerator.is_local_main_process:
+                    print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`")
+                test_mrpc(dispatch_batches, split_batches)
+                accelerator.state._reset_state()
+        print("test_gather_for_metrics_with_iterable_dataset")
+        test_gather_for_metrics_with_iterable_dataset()
+        print("test gather_for_metrics_with_non_tensor_objects_iterable_dataset")
+        test_gather_for_metrics_with_non_tensor_objects_iterable_dataset()
+
+    # MpDeviceLoader in TorchXLA is an asynchronous loader that preloads several batches into cache.
+    # This can cause the 'end_of_dataloader' of DataLoaderStateMixin to be set earlier than intended.
+    # Skip this test when TorchXLA is enabled.
+    if accelerator.state.distributed_type != DistributedType.XLA:
+        if accelerator.is_local_main_process:
+            print("**Test torch metrics**")
+        for split_batches in [True, False]:
+            for dispatch_batches in dispatch_batches_options:
+                dataloader_config = DataLoaderConfiguration(
+                    split_batches=split_batches, dispatch_batches=dispatch_batches
+                )
+                accelerator = Accelerator(dataloader_config=dataloader_config)
+                if accelerator.is_local_main_process:
+                    print(f"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99")
+                test_torch_metrics(accelerator, 99)
+                accelerator.state._reset_state()
+    if accelerator.is_local_main_process:
+        print("**Test last batch is not dropped when perfectly divisible**")
+    accelerator = Accelerator()
+    test_torch_metrics(accelerator, 512)
+    accelerator.state._reset_state()
+    if accelerator.is_local_main_process:
+        print("**Test that `drop_last` is taken into account**")
+    test_gather_for_metrics_drop_last()
+    accelerator.end_training()
+    accelerator.state._reset_state()
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py

@@ -0,0 +1,314 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import gc
+import json
+import os
+
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup, set_seed
+
+from accelerate import Accelerator, DistributedType
+from accelerate.utils import (
+    is_hpu_available,
+    is_mlu_available,
+    is_musa_available,
+    is_npu_available,
+    is_sdaa_available,
+    is_xpu_available,
+)
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+# Converting Bytes to Megabytes
+def b2mb(x):
+    return int(x / 2**20)
+
+
+# This context manager is used to track the peak memory usage of the process
+class TorchTracemalloc:
+    def __enter__(self):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.cuda.memory_allocated()
+        elif is_mlu_available():
+            torch.mlu.empty_cache()
+            torch.mlu.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.mlu.memory_allocated()
+        elif is_sdaa_available():
+            torch.sdaa.empty_cache()
+            torch.sdaa.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.sdaa.memory_allocated()
+        elif is_musa_available():
+            torch.musa.empty_cache()
+            torch.musa.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.musa.memory_allocated()
+        elif is_npu_available():
+            torch.npu.empty_cache()
+            torch.npu.reset_max_memory_allocated()  # reset the peak gauge to zero
+            self.begin = torch.npu.memory_allocated()
+        elif is_xpu_available():
+            torch.xpu.empty_cache()
+            torch.xpu.reset_peak_memory_stats()  # reset the peak gauge to zero
+            self.begin = torch.xpu.memory_allocated()
+        elif is_hpu_available():
+            # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process
+            torch.hpu.reset_peak_memory_stats()  # reset the peak gauge to zero
+            self.begin = torch.hpu.memory_allocated()
+        return self
+
+    def __exit__(self, *exc):
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+            self.end = torch.cuda.memory_allocated()
+            self.peak = torch.cuda.max_memory_allocated()
+        elif is_mlu_available():
+            torch.mlu.empty_cache()
+            self.end = torch.mlu.memory_allocated()
+            self.begin = torch.mlu.max_memory_allocated()
+        elif is_sdaa_available():
+            torch.sdaa.empty_cache()
+            self.end = torch.sdaa.memory_allocated()
+            self.begin = torch.sdaa.max_memory_allocated()
+        elif is_musa_available():
+            torch.musa.empty_cache()
+            self.end = torch.musa.memory_allocated()
+            self.begin = torch.musa.max_memory_allocated()
+        elif is_npu_available():
+            torch.npu.empty_cache()
+            self.end = torch.npu.memory_allocated()
+            self.peak = torch.npu.max_memory_allocated()
+        elif is_xpu_available():
+            torch.xpu.empty_cache()
+            self.end = torch.xpu.memory_allocated()
+            self.peak = torch.xpu.max_memory_allocated()
+        elif is_hpu_available():
+            # torch.hpu.empty_cache() # not available on hpu as it reserves all device memory for the current process
+            self.end = torch.hpu.memory_allocated()
+            self.peak = torch.hpu.max_memory_allocated()
+        self.used = b2mb(self.end - self.begin)
+        self.peaked = b2mb(self.peak - self.begin)
+        # print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")
+
+
+def get_dataloaders(
+    accelerator: Accelerator,
+    batch_size: int = 16,
+    model_name: str = "bert-base-cased",
+    n_train: int = 320,
+    n_val: int = 160,
+):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+            The name of the model to use.
+        n_train (`int`, *optional*):
+            The number of training examples to use.
+        n_val (`int`, *optional*):
+            The number of validation examples to use.
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    datasets = load_dataset(
+        "glue", "mrpc", split={"train": f"train[:{n_train}]", "validation": f"validation[:{n_val}]"}
+    )
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.XLA:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    # Initialize accelerator
+    accelerator = Accelerator()
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name, args.n_train, args.n_val)
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True)
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    if accelerator.state.deepspeed_plugin is not None:
+        gradient_accumulation_steps = accelerator.state.deepspeed_plugin.deepspeed_config[
+            "gradient_accumulation_steps"
+        ]
+    else:
+        gradient_accumulation_steps = 1
+    max_training_steps = (len(train_dataloader) * num_epochs) // gradient_accumulation_steps
+
+    # Instantiate scheduler
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We need to keep track of how many total steps we have iterated over
+    overall_step = 0
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Now we train the model
+    train_total_peak_memory = {}
+    for epoch in range(starting_epoch, num_epochs):
+        with TorchTracemalloc() as tracemalloc:
+            model.train()
+            for step, batch in enumerate(train_dataloader):
+                outputs = model(**batch)
+                loss = outputs.loss
+                loss = loss / gradient_accumulation_steps
+                accelerator.backward(loss)
+                if step % gradient_accumulation_steps == 0:
+                    optimizer.step()
+                    lr_scheduler.step()
+                    optimizer.zero_grad()
+
+                overall_step += 1
+
+        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
+        accelerator.print(f"Memory before entering the train : {b2mb(tracemalloc.begin)}")
+        accelerator.print(f"Memory consumed at the end of the train (end-begin): {tracemalloc.used}")
+        accelerator.print(f"Peak Memory consumed during the train (max-begin): {tracemalloc.peaked}")
+        accelerator.print(
+            f"Total Peak Memory consumed during the train (max): {tracemalloc.peaked + b2mb(tracemalloc.begin)}"
+        )
+        train_total_peak_memory[f"epoch-{epoch}"] = tracemalloc.peaked + b2mb(tracemalloc.begin)
+        if args.peak_memory_upper_bound is not None:
+            assert train_total_peak_memory[f"epoch-{epoch}"] <= args.peak_memory_upper_bound, (
+                "Peak memory usage exceeded the upper bound"
+            )
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        with open(os.path.join(args.output_dir, "peak_memory_utilization.json"), "w") as f:
+            json.dump(train_total_peak_memory, f)
+    accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--peak_memory_upper_bound",
+        type=float,
+        default=None,
+        help="The upper bound of peak memory usage in MB. If set, the training will throw an error if the peak memory usage exceeds this value.",
+    )
+    parser.add_argument(
+        "--n_train",
+        type=int,
+        default=320,
+        help="Number of training examples to use.",
+    )
+    parser.add_argument(
+        "--n_val",
+        type=int,
+        default=160,
+        help="Number of validation examples to use.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=1,
+        help="Number of train epochs.",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_performance.py

@@ -0,0 +1,299 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+import json
+import os
+from contextlib import nullcontext
+from pathlib import Path
+
+import evaluate
+import torch
+from datasets import load_dataset
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, get_linear_schedule_with_warmup
+
+from accelerate import Accelerator, DistributedType
+from accelerate.parallelism_config import ParallelismConfig
+from accelerate.utils import SAFE_WEIGHTS_NAME, set_seed
+from accelerate.utils.deepspeed import DummyOptim, DummyScheduler
+
+
+MAX_GPU_BATCH_SIZE = 16
+EVAL_BATCH_SIZE = 32
+
+
+def get_dataloaders(accelerator: Accelerator, batch_size: int = 16, model_name: str = "bert-base-cased"):
+    """
+    Creates a set of `DataLoader`s for the `glue` dataset.
+
+    Args:
+        accelerator (`Accelerator`):
+            An `Accelerator` object
+        batch_size (`int`, *optional*):
+            The batch size for the train and validation DataLoaders.
+        model_name (`str`, *optional*):
+    """
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    datasets = load_dataset("glue", "mrpc")
+
+    def tokenize_function(examples):
+        # max_length=None => use the model max length (it's actually the default)
+        outputs = tokenizer(examples["sentence1"], examples["sentence2"], truncation=True, max_length=None)
+        return outputs
+
+    # Apply the method we just defined to all the examples in all the splits of the dataset
+    tokenized_datasets = datasets.map(
+        tokenize_function, batched=True, remove_columns=["idx", "sentence1", "sentence2"], load_from_cache_file=False
+    )
+
+    # We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
+    # transformers library
+    tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+    def collate_fn(examples):
+        # On TPU it's best to pad everything to the same length or training will be very slow.
+        if accelerator.distributed_type == DistributedType.XLA:
+            return tokenizer.pad(examples, padding="max_length", max_length=128, return_tensors="pt")
+        return tokenizer.pad(examples, padding="longest", return_tensors="pt")
+
+    # Instantiate dataloaders.
+    train_dataloader = DataLoader(
+        tokenized_datasets["train"], shuffle=True, collate_fn=collate_fn, batch_size=batch_size
+    )
+    eval_dataloader = DataLoader(
+        tokenized_datasets["validation"], shuffle=False, collate_fn=collate_fn, batch_size=EVAL_BATCH_SIZE
+    )
+
+    return train_dataloader, eval_dataloader
+
+
+def training_function(config, args):
+    accelerator_kwargs = {}
+    # need this for DeepSpeed tests as `args.tp_size` would be None and `torch.distributed.init_device_mesh` would fail
+    if args.tp_size is not None:
+        accelerator_kwargs["parallelism_config"] = ParallelismConfig(tp_size=args.tp_size)
+
+    # Initialize accelerator
+    accelerator = Accelerator(**accelerator_kwargs)
+
+    # Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
+    lr = config["lr"]
+    num_epochs = int(config["num_epochs"])
+    seed = int(config["seed"])
+    batch_size = int(config["batch_size"])
+    model_name = args.model_name_or_path
+
+    set_seed(seed)
+    train_dataloader, eval_dataloader = get_dataloaders(accelerator, batch_size, model_name)
+
+    # Add TP related kwargs if provided
+    model_kwargs = {}
+    if args.tp_plan is not None:
+        model_kwargs["tp_plan"] = args.tp_plan
+    if args.tp_size is not None:
+        model_kwargs["tp_size"] = args.tp_size
+
+    # Instantiate the model (we build the model here so that the seed also control new weights initialization)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, return_dict=True, **model_kwargs)
+
+    if args.add_pad_token:
+        if model.config.pad_token_id is None:
+            model.config.pad_token_id = 0
+
+    # Instantiate optimizer
+    optimizer_cls = (
+        AdamW
+        if accelerator.state.deepspeed_plugin is None
+        or "optimizer" not in accelerator.state.deepspeed_plugin.deepspeed_config
+        else DummyOptim
+    )
+    optimizer = optimizer_cls(params=model.parameters(), lr=lr)
+
+    max_training_steps = len(train_dataloader) * num_epochs
+
+    # Instantiate scheduler
+    linear_decay_scheduler = False
+    if (
+        accelerator.state.deepspeed_plugin is None
+        or "scheduler" not in accelerator.state.deepspeed_plugin.deepspeed_config
+    ):
+        lr_scheduler = get_linear_schedule_with_warmup(
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=max_training_steps,
+        )
+        linear_decay_scheduler = True
+    else:
+        lr_scheduler = DummyScheduler(optimizer, total_num_steps=max_training_steps, warmup_num_steps=0)
+
+    # Prepare everything
+    # There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
+    # prepare method.
+    model, optimizer, train_dataloader, eval_dataloader, lr_scheduler = accelerator.prepare(
+        model, optimizer, train_dataloader, eval_dataloader, lr_scheduler
+    )
+
+    # We also need to keep track of the stating epoch so files are named properly
+    starting_epoch = 0
+
+    # Now we train the model
+    metric = evaluate.load("glue", "mrpc")
+    best_performance = 0
+    performance_metric = {}
+    expected_lr_after_first_optim_step = lr * (
+        1 - 1 / (max_training_steps / accelerator.num_processes / accelerator.gradient_accumulation_steps)
+    )
+    lr_scheduler_check_completed = False
+    for epoch in range(starting_epoch, num_epochs):
+        model.train()
+        for step, batch in enumerate(train_dataloader):
+            with accelerator.accumulate(model):
+                outputs = model(**batch)
+                loss = outputs.loss
+                accelerator.backward(loss)
+                context = nullcontext
+                if args.tp_plan is not None:
+                    from torch.distributed._tensor.experimental import implicit_replication
+
+                    context = implicit_replication
+                with context():
+                    optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
+                # assert the learning rate after first optimizer step
+                if (
+                    accelerator.sync_gradients
+                    and not lr_scheduler_check_completed
+                    and linear_decay_scheduler
+                    and accelerator.state.mixed_precision == "no"
+                ):
+                    assert lr_scheduler.get_last_lr()[0] == expected_lr_after_first_optim_step, (
+                        f"Wrong lr found at second step, expected {expected_lr_after_first_optim_step}, got {lr_scheduler.get_last_lr()[0]}"
+                    )
+                    lr_scheduler_check_completed = True
+
+        model.eval()
+        samples_seen = 0
+        for step, batch in enumerate(eval_dataloader):
+            # We could avoid this line since we set the accelerator with `device_placement=True`.
+            batch.to(accelerator.device)
+            with torch.no_grad():
+                outputs = model(**batch)
+            predictions = outputs.logits.argmax(dim=-1)
+            # It is slightly faster to call this once, than multiple times
+            predictions, references = accelerator.gather(
+                (predictions, batch["labels"])
+            )  # If we are in a multiprocess environment, the last batch has duplicates
+            if accelerator.use_distributed:
+                if step == len(eval_dataloader) - 1:
+                    predictions = predictions[: len(eval_dataloader.dataset) - samples_seen]
+                    references = references[: len(eval_dataloader.dataset) - samples_seen]
+                else:
+                    samples_seen += references.shape[0]
+            metric.add_batch(
+                predictions=predictions,
+                references=references,
+            )
+
+        eval_metric = metric.compute()
+        # Use accelerator.print to print only on the main process.
+        accelerator.print(f"epoch {epoch}:", eval_metric)
+        performance_metric[f"epoch-{epoch}"] = eval_metric["accuracy"]
+
+        if best_performance < eval_metric["accuracy"]:
+            best_performance = eval_metric["accuracy"]
+
+    # check that the LR is 0
+    if linear_decay_scheduler and accelerator.state.mixed_precision == "no":
+        assert lr_scheduler.get_last_lr()[0] == 0, (
+            f"Wrong lr found at last step, expected 0, got {lr_scheduler.get_last_lr()[0]}"
+        )
+
+    if args.performance_lower_bound is not None:
+        assert args.performance_lower_bound <= best_performance, (
+            f"Best performance metric {best_performance} is lower than the lower bound {args.performance_lower_bound}"
+        )
+
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        with open(os.path.join(args.output_dir, "all_results.json"), "w") as f:
+            json.dump(performance_metric, f)
+
+    # TODO: skip saving of the model test for TP until the feature lands
+    if args.tp_plan is None:
+        # Finally try saving the model
+        accelerator.save_model(model, args.output_dir)
+    accelerator.wait_for_everyone()
+    if args.tp_plan is None:
+        assert Path(args.output_dir, SAFE_WEIGHTS_NAME).exists(), (
+            "Model was not saved when calling `Accelerator.save_model`"
+        )
+    accelerator.end_training()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Simple example of training script tracking peak GPU memory usage.")
+    parser.add_argument(
+        "--model_name_or_path",
+        type=str,
+        default="bert-base-cased",
+        help="Path to pretrained model or model identifier from huggingface.co/models.",
+        required=False,
+    )
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=".",
+        help="Optional save directory where all checkpoint folders will be stored. Default is the current working directory.",
+    )
+    parser.add_argument(
+        "--performance_lower_bound",
+        type=float,
+        default=None,
+        help="Optional lower bound for the performance metric. If set, the training will throw error when the performance metric drops below this value.",
+    )
+    parser.add_argument(
+        "--num_epochs",
+        type=int,
+        default=3,
+        help="Number of train epochs.",
+    )
+    parser.add_argument(
+        "--add_pad_token",
+        type=bool,
+        default=False,
+        help="To add pad token if not exists.",
+    )
+    parser.add_argument(
+        "--tp_plan",
+        type=str,
+        default=None,
+        help="pass 'auto' to use TP",
+    )
+    parser.add_argument(
+        "--tp_size",
+        type=int,
+        default=None,
+        help="TP size to be used to shard the model",
+    )
+    args = parser.parse_args()
+    config = {"lr": 2e-5, "num_epochs": args.num_epochs, "seed": 42, "batch_size": 16}
+    training_function(config, args)
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_pippy.py

@@ -0,0 +1,117 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from transformers import (
+    BertConfig,
+    BertForMaskedLM,
+    GPT2Config,
+    GPT2ForSequenceClassification,
+)
+
+from accelerate import PartialState
+from accelerate.inference import prepare_pippy
+from accelerate.test_utils import torch_device
+from accelerate.utils import DistributedType, set_seed
+
+
+model_to_config = {
+    "bert": (BertForMaskedLM, BertConfig, 512),
+    "gpt2": (GPT2ForSequenceClassification, GPT2Config, 1024),
+}
+
+
+def get_model_and_data_for_text(model_name, device, num_processes: int = 2):
+    initializer, config, seq_len = model_to_config[model_name]
+    config_args = {}
+    # Eventually needed for batch inference tests on gpt-2 when bs != 1
+    # if model_name == "gpt2":
+    #     config_args["pad_token_id"] = 0
+    model_config = config(**config_args)
+    model = initializer(model_config)
+    kwargs = dict(low=0, high=model_config.vocab_size, device=device, dtype=torch.int64, requires_grad=False)
+    trace_input = torch.randint(size=(1, seq_len), **kwargs)
+    inference_inputs = torch.randint(size=(num_processes, seq_len), **kwargs)
+    return model, trace_input, inference_inputs
+
+
+def test_bert(batch_size: int = 2):
+    set_seed(42)
+    state = PartialState()
+    model, trace_input, inference_inputs = get_model_and_data_for_text("bert", "cpu", batch_size)
+    model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules)
+    # For inference args need to be a tuple
+    inputs = inference_inputs.to(torch_device)
+    with torch.no_grad():
+        output = model(inputs)
+    # Zach: Check that we just grab the real outputs we need at the end
+    if not state.is_last_process:
+        assert output is None, "Output was not generated on just the last process!"
+    else:
+        assert output is not None, "Output was not generated in the last process!"
+
+
+def test_gpt2(batch_size: int = 2):
+    set_seed(42)
+    state = PartialState()
+    model, trace_input, inference_inputs = get_model_and_data_for_text("gpt2", "cpu", batch_size)
+    model = prepare_pippy(model, example_args=(trace_input,), no_split_module_classes=model._no_split_modules)
+    # For inference args need to be a tuple
+    inputs = inference_inputs.to(torch_device)
+    with torch.no_grad():
+        output = model(inputs)
+    # Zach: Check that we just grab the real outputs we need at the end
+    if not state.is_last_process:
+        assert output is None, "Output was not generated on just the last process!"
+    else:
+        assert output is not None, "Output was not generated in the last process!"
+
+
+# Currently disabled, enable again once PyTorch pippy interface can trace a resnet34
+# def test_resnet(batch_size: int = 2):
+#     set_seed(42)
+#     state = PartialState()
+#     model = resnet34()
+#     input_tensor = torch.rand(1, 3, 224, 224)
+#     model = prepare_pippy(
+#         model,
+#         example_args=(input_tensor,),
+#     )
+#     inference_inputs = torch.rand(batch_size, 3, 224, 224)
+#     inputs = send_to_device(inference_inputs, torch_device)
+#     with torch.no_grad():
+#         output = model(inputs)
+#     # Zach: Check that we just grab the real outputs we need at the end
+#     if not state.is_last_process:
+#         assert output is None, "Output was not generated on just the last process!"
+#     else:
+#         assert output is not None, "Output was not generated in the last process!"
+
+
+if __name__ == "__main__":
+    state = PartialState()
+    state.print("Testing pippy integration...")
+    try:
+        if state.distributed_type in [DistributedType.MULTI_GPU, DistributedType.MULTI_XPU, DistributedType.MULTI_HPU]:
+            state.print("Testing GPT2...")
+            test_gpt2()
+            # Issue: When modifying the tokenizer for batch GPT2 inference, there's an issue
+            # due to references
+            # NameError: cannot access free variable 'chunk_args_list' where it is not associated with a value in enclosing scope
+            # test_gpt2(3)
+            state.print("Testing BERT...")
+            test_bert()
+        else:
+            print("Less than two GPUs found, not running tests!")
+    finally:
+        state.destroy_process_group()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/external_deps/test_zero3_integration.py

@@ -0,0 +1,59 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch.distributed
+
+from accelerate.test_utils import require_huggingface_suite, torch_device
+from accelerate.utils import is_transformers_available
+
+
+if is_transformers_available():
+    from transformers import AutoModel, TrainingArguments
+
+
+GPT2_TINY = "sshleifer/tiny-gpt2"
+
+
+@require_huggingface_suite
+def init_torch_dist_then_launch_deepspeed():
+    if torch_device == "xpu":
+        backend = "xccl"
+    elif torch_device == "hpu":
+        backend = "hccl"
+    else:
+        backend = "nccl"
+
+    torch.distributed.init_process_group(backend=backend)
+    deepspeed_config = {
+        "zero_optimization": {
+            "stage": 3,
+        },
+        "train_batch_size": "auto",
+        "train_micro_batch_size_per_gpu": "auto",
+    }
+    train_args = TrainingArguments(
+        output_dir="./",
+        deepspeed=deepspeed_config,
+    )
+    model = AutoModel.from_pretrained(GPT2_TINY)
+    assert train_args is not None
+    assert model is not None
+
+
+def main():
+    init_torch_dist_then_launch_deepspeed()
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/test_cli.py

@@ -0,0 +1,32 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+
+from accelerate.utils import is_xpu_available
+
+
+def main():
+    accelerator_type = "GPU"
+    num_accelerators = 0
+    if torch.cuda.is_available():
+        num_accelerators = torch.cuda.device_count()
+        accelerator_type = "GPU"
+    elif is_xpu_available():
+        num_accelerators = torch.xpu.device_count()
+        accelerator_type = "XPU"
+    print(f"Successfully ran on {num_accelerators} {accelerator_type}s")
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/accelerate/test_utils/scripts/test_ops.py

@@ -0,0 +1,181 @@
+#!/usr/bin/env python
+
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+
+from accelerate import PartialState
+from accelerate.test_utils.testing import assert_exception
+from accelerate.utils.dataclasses import DistributedType
+from accelerate.utils.operations import (
+    DistributedOperationException,
+    broadcast,
+    copy_tensor_to_devices,
+    gather,
+    gather_object,
+    pad_across_processes,
+    reduce,
+)
+
+
+def create_tensor(state):
+    return (torch.arange(state.num_processes) + 1.0 + (state.num_processes * state.process_index)).to(state.device)
+
+
+def test_gather(state):
+    tensor = create_tensor(state)
+    gathered_tensor = gather(tensor)
+    assert gathered_tensor.tolist() == list(range(1, state.num_processes**2 + 1))
+
+
+def test_gather_object(state):
+    # Gather objects in TorchXLA is not supported.
+    if state.distributed_type == DistributedType.XLA:
+        return
+    obj = [state.process_index]
+    gathered_obj = gather_object(obj)
+    assert len(gathered_obj) == state.num_processes, f"{gathered_obj}, {len(gathered_obj)} != {state.num_processes}"
+    assert gathered_obj == list(range(state.num_processes)), f"{gathered_obj} != {list(range(state.num_processes))}"
+
+
+def test_gather_non_contiguous(state):
+    # Skip this test because the 'is_contiguous' function of XLA tensor always returns True.
+    if state.distributed_type == DistributedType.XLA:
+        return
+
+    # Create a non-contiguous tensor (enforce non-contiguity after device memory allocation)
+    tensor = torch.arange(12, device=state.device).view(4, 3).t()
+    assert not tensor.is_contiguous()
+    # Shouldn't error out
+    _ = gather(tensor)
+
+
+def test_broadcast(state):
+    tensor = create_tensor(state)
+    broadcasted_tensor = broadcast(tensor)
+    assert broadcasted_tensor.shape == torch.Size([state.num_processes])
+    assert broadcasted_tensor.tolist() == list(range(1, state.num_processes + 1))
+
+
+def test_pad_across_processes(state):
+    # We need to pad the tensor with one more element if we are the main process
+    # to ensure that we can pad
+    if state.is_main_process:
+        tensor = torch.arange(state.num_processes + 1).to(state.device)
+    else:
+        tensor = torch.arange(state.num_processes).to(state.device)
+    padded_tensor = pad_across_processes(tensor)
+    assert padded_tensor.shape == torch.Size([state.num_processes + 1])
+    if not state.is_main_process:
+        assert padded_tensor.tolist() == list(range(0, state.num_processes)) + [0]
+
+
+def test_reduce_sum(state):
+    # For now runs on only two processes
+    if state.num_processes != 2:
+        return
+    tensor = create_tensor(state)
+    reduced_tensor = reduce(tensor, "sum")
+    truth_tensor = torch.tensor([4.0, 6]).to(state.device)
+    assert torch.allclose(reduced_tensor, truth_tensor), f"{reduced_tensor} != {truth_tensor}"
+
+
+def test_reduce_mean(state):
+    # For now runs on only two processes
+    if state.num_processes != 2:
+        return
+    tensor = create_tensor(state)
+    reduced_tensor = reduce(tensor, "mean")
+    truth_tensor = torch.tensor([2.0, 3]).to(state.device)
+    assert torch.allclose(reduced_tensor, truth_tensor), f"{reduced_tensor} != {truth_tensor}"
+
+
+def test_op_checker(state):
+    # Must be in a distributed state, and gathering is currently not supported in TorchXLA.
+    if state.distributed_type in [DistributedType.NO, DistributedType.XLA]:
+        return
+    state.debug = True
+    # `pad_across_processes`
+    if state.process_index == 0:
+        data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)}
+    else:
+        data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4, 5]]]).to(state.device)}
+
+    with assert_exception(DistributedOperationException):
+        pad_across_processes(data, dim=0)
+
+    # `reduce`
+    if state.process_index == 0:
+        data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)}
+    else:
+        data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]]).to(state.device)}
+
+    with assert_exception(DistributedOperationException):
+        reduce(data)
+
+    # `broadcast`
+    if state.process_index == 0:
+        data = {"tensor": torch.tensor([[0.0, 1, 2, 3, 4]]).to(state.device)}
+    else:
+        data = {"tensor": torch.tensor([[[0.0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]]).to(state.device)}
+
+    with assert_exception(DistributedOperationException):
+        broadcast(data)
+
+    state.debug = False
+
+
+def test_copy_tensor_to_devices(state):
+    if state.distributed_type not in [DistributedType.MULTI_GPU, DistributedType.XLA]:
+        return
+    if state.is_main_process:
+        tensor = torch.tensor([1, 2, 3], dtype=torch.int).to(state.device)
+    else:
+        tensor = None
+    tensor = copy_tensor_to_devices(tensor)
+    assert torch.allclose(tensor, torch.tensor([1, 2, 3], dtype=torch.int, device=state.device))
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+def main():
+    state = PartialState()
+    state.print(f"State: {state}")
+    state.print("testing gather")
+    test_gather(state)
+    state.print("testing gather_object")
+    test_gather_object(state)
+    state.print("testing gather non-contiguous")
+    test_gather_non_contiguous(state)
+    state.print("testing broadcast")
+    test_broadcast(state)
+    state.print("testing pad_across_processes")
+    test_pad_across_processes(state)
+    state.print("testing reduce_sum")
+    test_reduce_sum(state)
+    state.print("testing reduce_mean")
+    test_reduce_mean(state)
+    state.print("testing op_checker")
+    test_op_checker(state)
+    state.print("testing sending tensors across devices")
+    test_copy_tensor_to_devices(state)
+    state.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/dill/tests/__init__.py

@@ -0,0 +1,22 @@
+#!/usr/bin/env python
+#
+# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
+# Copyright (c) 2018-2024 The Uncertainty Quantification Foundation.
+# License: 3-clause BSD.  The full license text is available at:
+#  - https://github.com/uqfoundation/dill/blob/master/LICENSE
+"""
+to run this test suite, first build and install `dill`.
+
+  $ python -m pip install ../..
+
+
+then run the tests with:
+
+  $ python -m dill.tests
+
+
+or, if `nose` is installed:
+
+  $ nosetests  
+
+"""

Prism/LLaDA/LLaDA_Prism/.venv/lib/python3.12/site-packages/dill/tests/__main__.py

@@ -0,0 +1,35 @@
+#!/usr/bin/env python
+#
+# Author: Mike McKerns (mmckerns @caltech and @uqfoundation)
+# Copyright (c) 2018-2024 The Uncertainty Quantification Foundation.
+# License: 3-clause BSD.  The full license text is available at:
+#  - https://github.com/uqfoundation/dill/blob/master/LICENSE
+
+import glob
+import os
+import sys
+import subprocess as sp
+python = sys.executable
+try:
+    import pox
+    python = pox.which_python(version=True) or python
+except ImportError:
+    pass
+shell = sys.platform[:3] == 'win'
+
+suite = os.path.dirname(__file__) or os.path.curdir
+tests = glob.glob(suite + os.path.sep + 'test_*.py')
+
+
+if __name__ == '__main__':
+
+    failed = 0
+    for test in tests:
+        p = sp.Popen([python, test], shell=shell).wait()
+        if p:
+            print('F', end='', flush=True)
+            failed = 1
+        else:
+            print('.', end='', flush=True)
+    print('')
+    exit(failed)