"""simple docstring"""




import unittest

from transformers import AutoTokenizer, FalconConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
     import torch

     from transformers import (
         FalconForCausalLM,
         FalconForQuestionAnswering,
         FalconForSequenceClassification,
         FalconForTokenClassification,
         FalconModel,
     )
class 						__A							:






     def __init__(		self    ,					a__    ,					a__=3    ,					a__=7    ,					a__=True    ,					a__=True    ,					a__=False    ,					a__=True    ,					a__=99    ,					a__=32    ,					a__=5    ,					a__=4    ,					a__=37    ,					a__="gelu"    ,					a__=0.1    ,					a__=0.1    ,					a__=512    ,					a__=16    ,					a__=2    ,					a__=0.0_2    ,					a__=3    ,					a__=4    ,					a__=None    ,					):
       _lowerCAmelCase						:	Optional[int]        =					parent
       _lowerCAmelCase						:	List[Any]        =					batch_size
       _lowerCAmelCase						:	Any        =					seq_length
       _lowerCAmelCase						:	List[str]        =					is_training
       _lowerCAmelCase						:	str        =					use_input_mask
       _lowerCAmelCase						:	Union[str, Any]        =					use_token_type_ids
       _lowerCAmelCase						:	int        =					use_labels
       _lowerCAmelCase						:	Dict        =					vocab_size
       _lowerCAmelCase						:	str        =					hidden_size
       _lowerCAmelCase						:	List[str]        =					num_hidden_layers
       _lowerCAmelCase						:	Any        =					num_attention_heads
       _lowerCAmelCase						:	Optional[Any]        =					intermediate_size
       _lowerCAmelCase						:	Optional[int]        =					hidden_act
       _lowerCAmelCase						:	Dict        =					hidden_dropout_prob
       _lowerCAmelCase						:	Dict        =					attention_probs_dropout_prob
       _lowerCAmelCase						:	Optional[Any]        =					max_position_embeddings
       _lowerCAmelCase						:	Tuple        =					type_vocab_size
       _lowerCAmelCase						:	Optional[Any]        =					type_sequence_label_size
       _lowerCAmelCase						:	List[Any]        =					initializer_range
       _lowerCAmelCase						:	Union[str, Any]        =					num_labels
       _lowerCAmelCase						:	Union[str, Any]        =					num_choices
       _lowerCAmelCase						:	Any        =					scope






     def    __A  (		self  ):
       _lowerCAmelCase						:	Dict        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.vocab_size  )

       _lowerCAmelCase						:	Optional[int]        =					None
       if self.use_input_mask:
         _lowerCAmelCase						:	Optional[int]        =					random_attention_mask([self.batch_size, self.seq_length]  )

       _lowerCAmelCase						:	Tuple        =					None

       _lowerCAmelCase						:	List[str]        =					None
       _lowerCAmelCase						:	Any        =					None
       _lowerCAmelCase						:	Any        =					None
       if self.use_labels:
         _lowerCAmelCase						:	str        =					ids_tensor([self.batch_size]    ,					self.type_sequence_label_size  )
         _lowerCAmelCase						:	Any        =					ids_tensor([self.batch_size, self.seq_length]    ,					self.num_labels  )
         _lowerCAmelCase						:	Any        =					ids_tensor([self.batch_size]    ,					self.num_choices  )

       _lowerCAmelCase						:	List[Any]        =					self.get_config()

       return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels






     def    __A  (		self  ):
       return FalconConfig(
           vocab_size=self.vocab_size    ,					hidden_size=self.hidden_size    ,					num_hidden_layers=self.num_hidden_layers    ,					num_attention_heads=self.num_attention_heads    ,					intermediate_size=self.intermediate_size    ,					hidden_act=self.hidden_act    ,					hidden_dropout_prob=self.hidden_dropout_prob    ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob    ,					max_position_embeddings=self.max_position_embeddings    ,					type_vocab_size=self.type_vocab_size    ,					is_decoder=a__    ,					initializer_range=self.initializer_range    ,					pad_token_id=1    ,					new_decoder_architecture=a__    ,					)






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__  ):
       _lowerCAmelCase						:	Union[str, Any]        =					FalconModel(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	int        =					model(a__    ,					attention_mask=a__  )
       _lowerCAmelCase						:	Dict        =					model(a__  )
       self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Optional[int]        =					True
       _lowerCAmelCase						:	Tuple        =					FalconModel(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[Any]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					)
       _lowerCAmelCase						:	int        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					)
       _lowerCAmelCase						:	Optional[int]        =					model(a__    ,					attention_mask=a__  )
       self.parent.assertEqual(result.last_hidden_state.shape    ,					(self.batch_size, self.seq_length, self.hidden_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Dict        =					FalconForCausalLM(config=a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.parent.assertEqual(result.logits.shape    ,					(self.batch_size, self.seq_length, self.vocab_size)  )






     def    __A  (		self    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					a__    ,					):
       _lowerCAmelCase						:	Optional[int]        =					True
       _lowerCAmelCase						:	List[str]        =					True
       _lowerCAmelCase						:	List[Any]        =					FalconForCausalLM(config=a__  )
       model.to(a__  )
       model.eval()

       # first forward pass
       _lowerCAmelCase						:	Dict        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					use_cache=a__    ,					)
       _lowerCAmelCase						:	Any        =					outputs.past_key_values

       # create hypothetical multiple next token and extent to next_input_ids
       _lowerCAmelCase						:	Any        =					ids_tensor((self.batch_size, 3)    ,					config.vocab_size  )
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor((self.batch_size, 3)    ,					vocab_size=2  )

       # append to next input_ids and
       _lowerCAmelCase						:	Any        =					torch.cat([input_ids, next_tokens]    ,					dim=-1  )
       _lowerCAmelCase						:	str        =					torch.cat([input_mask, next_mask]    ,					dim=-1  )

       _lowerCAmelCase						:	List[str]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					output_hidden_states=a__    ,					)["""hidden_states"""][0]
       _lowerCAmelCase						:	Optional[Any]        =					model(
           a__    ,					attention_mask=a__    ,					encoder_hidden_states=a__    ,					encoder_attention_mask=a__    ,					past_key_values=a__    ,					output_hidden_states=a__    ,					)["""hidden_states"""][0]

       # select random slice
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor((1,)    ,					output_from_past.shape[-1]  ).item()
       _lowerCAmelCase						:	Tuple        =					output_from_no_past[:, -3:, random_slice_idx].detach()
       _lowerCAmelCase						:	Union[str, Any]        =					output_from_past[:, :, random_slice_idx].detach()

       self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]  )

       # test that outputs are equal for slice
       self.parent.assertTrue(torch.allclose(a__    ,					a__    ,					atol=1e-3  )  )






     def    __A  (		self  ):
       _lowerCAmelCase						:	Union[str, Any]        =					self.prepare_config_and_inputs()
       (
           (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     (
           _lowerCAmelCase
       )      ,     
       )						:	Tuple        =					config_and_inputs
       _lowerCAmelCase						:	Optional[Any]        =					{"""input_ids""": input_ids, """attention_mask""": input_mask}
       return config, inputs_dict







@require_torch
class 						__A							(			SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     SCREAMING_SNAKE_CASE_   ,     unittest.TestCase		):
     _UpperCamelCase							:      Optional[int]		 =		(
         (
             FalconModel,
             FalconForCausalLM,
             FalconForSequenceClassification,
             FalconForTokenClassification,
             FalconForQuestionAnswering,
         )
         if is_torch_available()
         else ()
     )
     _UpperCamelCase							:      int		 =		(FalconForCausalLM,) if is_torch_available() else ()
     _UpperCamelCase							:      Any		 =		(
         {
             "feature-extraction": FalconModel,
             "text-classification": FalconForSequenceClassification,
             "text-generation": FalconForCausalLM,
             "question-answering": FalconForQuestionAnswering,
             "token-classification": FalconForTokenClassification,
             "zero-shot": FalconForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     _UpperCamelCase							:      Tuple		 =		False
     _UpperCamelCase							:      Optional[Any]		 =		False






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[int]        =					FalconModelTester(self  )
       _lowerCAmelCase						:	Dict        =					ConfigTester(self    ,					config_class=a__    ,					hidden_size=37  )






     def    __A  (		self  ):
       self.config_tester.run_common_tests()






     def    __A  (		self  ):
       _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     *_lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs()
       for alibi in [True, False]:
         _lowerCAmelCase						:	Tuple        =					alibi
         self.model_tester.create_and_check_model(a__    ,					*a__  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase						:	Optional[Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	List[Any]        =					3
       _lowerCAmelCase						:	Union[str, Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Dict        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	Optional[Any]        =					ids_tensor([self.model_tester.batch_size]    ,					self.model_tester.type_sequence_label_size  )
       _lowerCAmelCase						:	Optional[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase						:	Tuple        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	Optional[Any]        =					3
       _lowerCAmelCase						:	Optional[int]        =					"""single_label_classification"""
       _lowerCAmelCase						:	List[Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Dict        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	List[str]        =					ids_tensor([self.model_tester.batch_size]    ,					self.model_tester.type_sequence_label_size  )
       _lowerCAmelCase						:	Optional[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[str]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	List[str]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	List[Any]        =					FalconForCausalLM(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	List[str]        =					model(a__    ,					use_cache=a__  )
       _lowerCAmelCase						:	Dict        =					input_ids.shape[0]
       _lowerCAmelCase						:	Dict        =					model._convert_to_rw_cache(result.past_key_values  )
       _lowerCAmelCase						:	int        =					model._convert_cache_to_standard_format(a__    ,					a__  )
       for layer in range(len(a__  )  ):
         for tensor_idx in range(2  ):
           self.assertTrue(rw_cache[layer][tensor_idx].ndim == 3  )
           self.assertTrue(result.past_key_values[layer][tensor_idx].ndim == 4  )
           self.assertTrue(
               torch.all(result.past_key_values[layer][tensor_idx] == standard_cache[layer][tensor_idx]  )  )






     def    __A  (		self  ):
       _lowerCAmelCase      ,     _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()
       _lowerCAmelCase						:	str        =					3
       _lowerCAmelCase						:	List[Any]        =					"""multi_label_classification"""
       _lowerCAmelCase						:	List[Any]        =					input_dict["""input_ids"""]
       _lowerCAmelCase						:	Tuple        =					input_ids.ne(1  ).to(a__  )
       _lowerCAmelCase						:	Optional[int]        =					ids_tensor(
           [self.model_tester.batch_size, config.num_labels]    ,					self.model_tester.type_sequence_label_size  ).to(torch.float  )
       _lowerCAmelCase						:	List[Any]        =					FalconForSequenceClassification(a__  )
       model.to(a__  )
       model.eval()
       _lowerCAmelCase						:	Union[str, Any]        =					model(a__    ,					attention_mask=a__    ,					labels=a__  )
       self.assertEqual(result.logits.shape    ,					(self.model_tester.batch_size, self.model_tester.num_labels)  )






     def    __A  (		self  ):
       # Falcon can have different numbers of KV-heads than the number of query heads, so we need
       # to override this test to use the right head counts.
       for model_class in self.all_generative_model_classes:
         _lowerCAmelCase      ,     _lowerCAmelCase						:	Union[str, Any]        =					self.model_tester.prepare_config_and_inputs_for_common()

         # If it doesn't support cache, pass the test
         if not hasattr(a__    ,					"""use_cache"""  ):
           return

         _lowerCAmelCase						:	Dict        =					model_class(a__  ).to(a__  )
         if "use_cache" not in inputs:
           _lowerCAmelCase						:	Union[str, Any]        =					True
         _lowerCAmelCase						:	int        =					model(**a__  )

         # If "past_key_values" is not returned, pass the test (e.g. RWKV uses a different cache name and format)
         if "past_key_values" not in outputs:
           return

         _lowerCAmelCase						:	int        =					(
             getattr(a__    ,					"""decoder_layers"""    ,					a__  )
             or getattr(a__    ,					"""num_decoder_layers"""    ,					a__  )
             or config.num_hidden_layers
         )
         _lowerCAmelCase						:	Dict        =					getattr(a__    ,					"""num_kv_heads"""    ,					config.num_attention_heads  )
         _lowerCAmelCase						:	Optional[Any]        =					getattr(a__    ,					"""d_model"""    ,					config.hidden_size  )
         _lowerCAmelCase						:	List[Any]        =					embed_dim // num_attention_heads

         _lowerCAmelCase						:	int        =					outputs["""past_key_values"""]
         self.assertEqual(len(a__  )    ,					a__  )

         _lowerCAmelCase      ,     _lowerCAmelCase						:	Tuple        =					inputs["""input_ids"""].shape
         for i in range(a__  ):
           if config.new_decoder_architecture:
             _lowerCAmelCase						:	Any        =					config.num_attention_heads
           elif config.multi_query:
             _lowerCAmelCase						:	Dict        =					1
           self.assertEqual(len(past_kv[0]  )    ,					2  )  # K V for the decoder = 2
           self.assertEqual(
               past_kv[i][0].shape    ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)  )
           self.assertEqual(
               past_kv[i][1].shape    ,					(batch_size, num_attention_heads, seq_length, per_head_embed_dim)  )







@require_torch
class 						__A							(			unittest.TestCase		):






     @slow
     def    __A  (		self  ):
       _lowerCAmelCase						:	Any        =					AutoTokenizer.from_pretrained("""Rocketknight1/falcon-rw-1b"""  )
       _lowerCAmelCase						:	Dict        =					FalconForCausalLM.from_pretrained("""Rocketknight1/falcon-rw-1b"""  )
       model.eval()
       model.to(a__  )
       _lowerCAmelCase						:	List[Any]        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

       _lowerCAmelCase						:	Dict        =					(
           """My favorite food is pizza. I love it so much that I have a pizza party every year for my birthday."""
       )

       _lowerCAmelCase						:	List[Any]        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=19  )
       _lowerCAmelCase						:	int        =					tokenizer.batch_decode(a__  )[0]

       self.assertEqual(a__    ,					a__  )






     @slow
     def    __A  (		self  ):
       # The big models are way too big for the CI, so we use tiny random models that resemble their
       # architectures but with much smaller and fewer layers
       for repo in ["Rocketknight1/tiny-random-falcon-7b", "Rocketknight1/tiny-random-falcon-40b"]:
         _lowerCAmelCase						:	List[Any]        =					AutoTokenizer.from_pretrained(a__  )
         _lowerCAmelCase						:	Any        =					FalconForCausalLM.from_pretrained(a__  )
         model.eval()
         model.to(a__  )
         _lowerCAmelCase						:	Union[str, Any]        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

         # We just test that these run without errors - the models are randomly initialized
         # and so the actual text outputs will be garbage
         model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=4  )
         model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=4  )
         model.generate(**a__    ,					num_beams=2    ,					max_new_tokens=4  )






     @slow
     def    __A  (		self  ):
       # The big models are way too big for the CI, so we use tiny random models that resemble their
       # architectures but with much smaller and fewer layers
       with torch.no_grad():
         for repo in [
             "Rocketknight1/falcon-rw-1b",
             "Rocketknight1/tiny-random-falcon-7b",
             "Rocketknight1/tiny-random-falcon-40b",
         ]:
           _lowerCAmelCase						:	int        =					AutoTokenizer.from_pretrained(a__  )
           _lowerCAmelCase						:	Union[str, Any]        =					FalconForCausalLM.from_pretrained(a__  )
           model.eval()
           model.to(device=a__  )
           _lowerCAmelCase						:	int        =					tokenizer("""My favorite food is"""    ,					return_tensors="""pt"""  ).to(a__  )

           # Test results are the same with and without cache
           _lowerCAmelCase						:	str        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=20    ,					use_cache=a__  )
           _lowerCAmelCase						:	str        =					model.generate(**a__    ,					do_sample=a__    ,					max_new_tokens=20    ,					use_cache=a__  )
           self.assertTrue((outputs_cache - outputs_no_cache).sum().item() == 0  )





"""simple docstring"""




from .constants import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SAFE_WEIGHTS_INDEX_NAME,
    SAFE_WEIGHTS_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    TORCH_LAUNCH_PARAMS,
    WEIGHTS_INDEX_NAME,
    WEIGHTS_NAME,
)
from .dataclasses import (
    BnbQuantizationConfig,
    ComputeEnvironment,
    CustomDtype,
    DeepSpeedPlugin,
    DistributedDataParallelKwargs,
    DistributedType,
    DynamoBackend,
    FPaRecipeKwargs,
    FullyShardedDataParallelPlugin,
    GradientAccumulationPlugin,
    GradScalerKwargs,
    InitProcessGroupKwargs,
    KwargsHandler,
    LoggerType,
    MegatronLMPlugin,
    PrecisionType,
    ProjectConfiguration,
    RNGType,
    SageMakerDistributedType,
    TensorInformation,
    TorchDynamoPlugin,
)
from .environment import get_int_from_env, parse_choice_from_env, parse_flag_from_env
from .imports import (
    get_ccl_version,
    is_abit_bnb_available,
    is_abit_bnb_available,
    is_aim_available,
    is_bfaa_available,
    is_bnb_available,
    is_botoa_available,
    is_ccl_available,
    is_comet_ml_available,
    is_datasets_available,
    is_deepspeed_available,
    is_fpa_available,
    is_ipex_available,
    is_megatron_lm_available,
    is_mlflow_available,
    is_mps_available,
    is_npu_available,
    is_rich_available,
    is_safetensors_available,
    is_sagemaker_available,
    is_tensorboard_available,
    is_tpu_available,
    is_transformers_available,
    is_wandb_available,
    is_xpu_available,
)
from .modeling import (
    check_device_map,
    check_tied_parameters_in_config,
    check_tied_parameters_on_same_device,
    compute_module_sizes,
    convert_file_size_to_int,
    dtype_byte_size,
    find_tied_parameters,
    get_balanced_memory,
    get_max_layer_size,
    get_max_memory,
    get_mixed_precision_context_manager,
    id_tensor_storage,
    infer_auto_device_map,
    load_checkpoint_in_model,
    load_offloaded_weights,
    load_state_dict,
    named_module_tensors,
    retie_parameters,
    set_module_tensor_to_device,
    shard_checkpoint,
)
from .offload import (
    OffloadedWeightsLoader,
    PrefixedDataset,
    extract_submodules_state_dict,
    load_offloaded_weight,
    offload_state_dict,
    offload_weight,
    save_offload_index,
)
from .operations import (
    broadcast,
    broadcast_object_list,
    concatenate,
    convert_outputs_to_fpaa,
    convert_to_fpaa,
    find_batch_size,
    find_device,
    gather,
    gather_object,
    get_data_structure,
    honor_type,
    initialize_tensors,
    is_namedtuple,
    is_tensor_information,
    is_torch_tensor,
    listify,
    pad_across_processes,
    recursively_apply,
    reduce,
    send_to_device,
    slice_tensors,
)
from .versions import compare_versions, is_torch_version


if is_deepspeed_available():
     from .deepspeed import (
         DeepSpeedEngineWrapper,
         DeepSpeedOptimizerWrapper,
         DeepSpeedSchedulerWrapper,
         DummyOptim,
         DummyScheduler,
         HfDeepSpeedConfig,
     )

from .bnb import has_abit_bnb_layers, load_and_quantize_model
from .fsdp_utils import load_fsdp_model, load_fsdp_optimizer, save_fsdp_model, save_fsdp_optimizer
from .launch import (
    PrepareForLaunch,
    _filter_args,
    prepare_deepspeed_cmd_env,
    prepare_multi_gpu_env,
    prepare_sagemager_args_inputs,
    prepare_simple_launcher_cmd_env,
    prepare_tpu,
)
from .megatron_lm import (
    AbstractTrainStep,
    BertTrainStep,
    GPTTrainStep,
    MegatronEngine,
    MegatronLMDummyDataLoader,
    MegatronLMDummyScheduler,
    MegatronLMOptimizerWrapper,
    MegatronLMSchedulerWrapper,
    TaTrainStep,
    avg_losses_across_data_parallel_group,
    gather_across_data_parallel_groups,
)
from .megatron_lm import initialize as megatron_lm_initialize
from .megatron_lm import prepare_data_loader as megatron_lm_prepare_data_loader
from .megatron_lm import prepare_model as megatron_lm_prepare_model
from .megatron_lm import prepare_optimizer as megatron_lm_prepare_optimizer
from .megatron_lm import prepare_scheduler as megatron_lm_prepare_scheduler
from .memory import find_executable_batch_size, release_memory
from .other import (
    extract_model_from_parallel,
    get_pretty_name,
    is_port_in_use,
    merge_dicts,
    patch_environment,
    save,
    wait_for_everyone,
    write_basic_config,
)
from .random import set_seed, synchronize_rng_state, synchronize_rng_states
from .torch_xla import install_xla
from .tqdm import tqdm
from .transformer_engine import convert_model, has_transformer_engine_layers




