Skip to content

feature: add mps kernel for compute_puff_advantage #422

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
hysmio wants to merge 8 commits into
base: 3.0
Choose a base branch
from
Open

feature: add mps kernel for compute_puff_advantage #422

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
1f02571
feature: add mps kernel
hysmio Nov 18, 2025
12ccbd1
fix: initialisation weirdness
hysmio Nov 22, 2025
31a9543
feature: add test & benchmarks at varying steps/horizon
hysmio Nov 22, 2025
cb781d3
fix: test copy & paste
hysmio Nov 22, 2025
b3445cd
fix: use in pufferrl, missed in merge conflict
hysmio Nov 22, 2025
a51c22c
fix: use accelerator api instead of cuda directly
hysmio Nov 23, 2025
e112b3f
fix: test_mps_advantage check diff in benchmarks
hysmio Nov 24, 2025
c6f4a98
feature: use accelerator api more generally
hysmio Nov 24, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion pufferlib/config/default.ini
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ project = ablations
seed = 42
torch_deterministic = True
cpu_offload = False
device = cuda
device = default
Copy link
Author

@hysmio hysmio Nov 24, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure if this is controversial or not, but just means people don't need to manually pass --train.device mps every time, automatically select the best one based off of the get_accelerator function

optimizer = muon
precision = float32
total_timesteps = 10_000_000
Expand Down
11 changes: 9 additions & 2 deletions pufferlib/pufferl.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -694,6 +694,12 @@ def abbreviate(num, b2, c2):
else:
return f'{b2}{num/1e12:.2f}{c2}T'

def get_accelerator(device):
if device == 'default':
return torch.accelerator.current_accelerator() if torch.accelerator.is_available() else 'cpu'
else:
return device

def duration(seconds, b2, c2):
if seconds < 0:
return f"{b2}0{c2}s"
Expand Down Expand Up @@ -994,7 +1000,7 @@ def eval(env_name, args=None, vecenv=None, policy=None):
ob, info = vecenv.reset()
driver = vecenv.driver_env
num_agents = vecenv.observation_space.shape[0]
device = args['train']['device']
device = get_accelerator(args['train']['device'])

state = {}
if args['train']['use_rnn']:
Expand Down Expand Up @@ -1146,7 +1152,7 @@ def load_policy(args, vecenv, env_name=''):
module_name = 'pufferlib.ocean' if package == 'ocean' else f'pufferlib.environments.{package}'
env_module = importlib.import_module(module_name)

device = args['train']['device']
device = get_accelerator(args['train']['device'])
policy_cls = getattr(env_module.torch, args['policy_name'])
policy = policy_cls(vecenv.driver_env, **args['policy'])

Expand Down Expand Up @@ -1284,6 +1290,7 @@ def auto_type(value):

args['train']['env'] = args['env_name'] or '' # for trainer dashboard
args['train']['use_rnn'] = args['rnn_name'] is not None
args['train']['device'] = get_accelerator(args['train']['device'])
return args

def main():
Expand Down